kiln-ai 0.20.1__py3-none-any.whl → 0.21.0__py3-none-any.whl

kiln_ai/adapters/eval/eval_runner.py

@@ -160,8 +160,12 @@ class EvalRunner:
         """
         jobs = self.collect_tasks()
 
-        runner = AsyncJobRunner(concurrency=concurrency)
-        async for progress in runner.run(jobs, self.run_job):
+        runner = AsyncJobRunner(
+            concurrency=concurrency,
+            jobs=jobs,
+            run_job_fn=self.run_job,
+        )
+        async for progress in runner.run():
             yield progress
 
     async def run_job(self, job: EvalJob) -> bool:

kiln_ai/adapters/eval/test_base_eval.py

@@ -307,9 +307,7 @@ async def test_run_method():
     evaluator = EvalTester(eval_config, run_config.run_config())
 
     # Run the evaluation
-    task_run, eval_scores, intermediate_outputs = await evaluator.run_task_and_eval(
-        "test input"
-    )
+    task_run, eval_scores, _ = await evaluator.run_task_and_eval("test input")
 
     # Verify task run was created
     assert task_run.input == "test input"

kiln_ai/adapters/eval/test_g_eval.py

@@ -188,7 +188,7 @@ async def test_run_g_eval_e2e(
     g_eval = GEval(test_eval_config, test_run_config)
 
     # Run the evaluation
-    task_run, scores, intermediate_outputs = await g_eval.run_task_and_eval("chickens")
+    _, scores, intermediate_outputs = await g_eval.run_task_and_eval("chickens")
 
     # Verify the evaluation results
     assert isinstance(scores, dict)

kiln_ai/adapters/extractors/__init__.py

@@ -0,0 +1,18 @@
+"""
+File extractors for processing different document types.
+
+This package provides a framework for extracting content from files
+using different extraction methods.
+"""
+
+from . import base_extractor, extractor_registry, extractor_runner, litellm_extractor
+from .base_extractor import ExtractionInput, ExtractionOutput
+
+__all__ = [
+    "ExtractionInput",
+    "ExtractionOutput",
+    "base_extractor",
+    "extractor_registry",
+    "extractor_runner",
+    "litellm_extractor",
+]

kiln_ai/adapters/extractors/base_extractor.py

@@ -0,0 +1,72 @@
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+from kiln_ai.datamodel.extraction import ExtractorConfig, OutputFormat
+
+logger = logging.getLogger(__name__)
+
+
+class ExtractionInput(BaseModel):
+    path: Path | str = Field(description="The absolute path to the file to extract.")
+    mime_type: str = Field(description="The mime type of the file.")
+
+
+class ExtractionOutput(BaseModel):
+    """
+    The output of an extraction. This is the data that will be saved to the data store.
+    """
+
+    is_passthrough: bool = Field(
+        default=False, description="Whether the extractor returned the file as is."
+    )
+    content_format: OutputFormat = Field(
+        description="The format of the extracted data."
+    )
+    content: str = Field(description="The extracted data.")
+
+
+class BaseExtractor(ABC):
+    """
+    Base class for all extractors.
+
+    Should be subclassed by each extractor.
+    """
+
+    def __init__(self, extractor_config: ExtractorConfig):
+        self.extractor_config = extractor_config
+
+    @abstractmethod
+    async def _extract(self, extraction_input: ExtractionInput) -> ExtractionOutput:
+        pass
+
+    async def extract(
+        self,
+        extraction_input: ExtractionInput,
+    ) -> ExtractionOutput:
+        """
+        Extracts content from a file by delegating to the concrete extractor implementation.
+        """
+        try:
+            if self._should_passthrough(extraction_input.mime_type):
+                return ExtractionOutput(
+                    is_passthrough=True,
+                    content=Path(extraction_input.path).read_text(encoding="utf-8"),
+                    content_format=self.extractor_config.output_format,
+                )
+
+            return await self._extract(
+                extraction_input,
+            )
+        except Exception as e:
+            raise ValueError(f"Error extracting {extraction_input.path}: {e}") from e
+
+    def _should_passthrough(self, mime_type: str) -> bool:
+        return mime_type.lower() in {
+            mt.lower() for mt in self.extractor_config.passthrough_mimetypes
+        }
+
+    def output_format(self) -> OutputFormat:
+        return self.extractor_config.output_format

kiln_ai/adapters/extractors/encoding.py

@@ -0,0 +1,20 @@
+import base64
+
+
+def to_base64_url(mime_type: str, bytes: bytes) -> str:
+    base64_url = f"data:{mime_type};base64,{base64.b64encode(bytes).decode('utf-8')}"
+    return base64_url
+
+
+def from_base64_url(base64_url: str) -> bytes:
+    if not base64_url.startswith("data:") or "," not in base64_url:
+        raise ValueError("Invalid base64 URL format")
+
+    parts = base64_url.split(",")
+    if len(parts) != 2:
+        raise ValueError("Invalid base64 URL format")
+
+    try:
+        return base64.b64decode(parts[1])
+    except Exception as e:
+        raise ValueError(f"Failed to decode base64 data: {e}")

kiln_ai/adapters/extractors/extractor_registry.py

@@ -0,0 +1,44 @@
+from kiln_ai.adapters.extractors.base_extractor import BaseExtractor
+from kiln_ai.adapters.extractors.litellm_extractor import LitellmExtractor
+from kiln_ai.adapters.ml_model_list import ModelProviderName
+from kiln_ai.adapters.provider_tools import (
+    core_provider,
+    lite_llm_core_config_for_provider,
+)
+from kiln_ai.datamodel.extraction import ExtractorConfig, ExtractorType
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+from kiln_ai.utils.filesystem_cache import FilesystemCache
+
+
+def extractor_adapter_from_type(
+    extractor_type: ExtractorType,
+    extractor_config: ExtractorConfig,
+    filesystem_cache: FilesystemCache | None = None,
+) -> BaseExtractor:
+    match extractor_type:
+        case ExtractorType.LITELLM:
+            try:
+                provider_enum = ModelProviderName(extractor_config.model_provider_name)
+            except ValueError:
+                raise ValueError(
+                    f"Unsupported model provider name: {extractor_config.model_provider_name}. "
+                )
+
+            core_provider_name = core_provider(
+                extractor_config.model_name, provider_enum
+            )
+
+            provider_config = lite_llm_core_config_for_provider(core_provider_name)
+            if provider_config is None:
+                raise ValueError(
+                    f"No configuration found for core provider: {core_provider_name.value}. "
+                )
+
+            return LitellmExtractor(
+                extractor_config,
+                provider_config,
+                filesystem_cache,
+            )
+        case _:
+            # type checking will catch missing cases
+            raise_exhaustive_enum_error(extractor_type)

kiln_ai/adapters/extractors/extractor_runner.py

@@ -0,0 +1,112 @@
+import logging
+from collections import defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import AsyncGenerator, Dict, List, Set
+
+from kiln_ai.adapters.extractors.base_extractor import BaseExtractor, ExtractionInput
+from kiln_ai.adapters.extractors.extractor_registry import extractor_adapter_from_type
+from kiln_ai.datamodel.basemodel import ID_TYPE, KilnAttachmentModel
+from kiln_ai.datamodel.extraction import (
+    Document,
+    Extraction,
+    ExtractionSource,
+    ExtractorConfig,
+)
+from kiln_ai.utils.async_job_runner import AsyncJobRunner, Progress
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ExtractorJob:
+    doc: Document
+    extractor_config: ExtractorConfig
+
+
+class ExtractorRunner:
+    def __init__(
+        self,
+        documents: List[Document],
+        extractor_configs: List[ExtractorConfig],
+    ):
+        if len(extractor_configs) == 0:
+            raise ValueError("Extractor runner requires at least one extractor config")
+
+        self.documents = documents
+        self.extractor_configs = extractor_configs
+
+    def collect_jobs(self) -> List[ExtractorJob]:
+        jobs = []
+
+        # we want to avoid re-running the same document for the same extractor config
+        already_extracted: Dict[ID_TYPE, Set[ID_TYPE]] = defaultdict(set)
+        for document in self.documents:
+            for extraction in document.extractions():
+                already_extracted[extraction.extractor_config_id].add(document.id)
+
+        for extractor_config in self.extractor_configs:
+            for document in self.documents:
+                if document.id not in already_extracted[extractor_config.id]:
+                    jobs.append(
+                        ExtractorJob(
+                            doc=document,
+                            extractor_config=extractor_config,
+                        )
+                    )
+
+        return jobs
+
+    async def run(self, concurrency: int = 25) -> AsyncGenerator[Progress, None]:
+        jobs = self.collect_jobs()
+
+        runner = AsyncJobRunner(
+            concurrency=concurrency,
+            jobs=jobs,
+            run_job_fn=self.run_job,
+        )
+        async for progress in runner.run():
+            yield progress
+
+    async def run_job(self, job: ExtractorJob) -> bool:
+        try:
+            extractor = extractor_adapter_from_type(
+                job.extractor_config.extractor_type,
+                job.extractor_config,
+            )
+            if not isinstance(extractor, BaseExtractor):
+                raise ValueError("Not able to create extractor from extractor config")
+
+            if job.doc.path is None:
+                raise ValueError("Document path is not set")
+
+            output = await extractor.extract(
+                extraction_input=ExtractionInput(
+                    path=Path(
+                        job.doc.original_file.attachment.resolve_path(
+                            job.doc.path.parent
+                        )
+                    ),
+                    mime_type=job.doc.original_file.mime_type,
+                )
+            )
+
+            extraction = Extraction(
+                parent=job.doc,
+                extractor_config_id=job.extractor_config.id,
+                output=KilnAttachmentModel.from_data(
+                    data=output.content,
+                    mime_type=output.content_format,
+                ),
+                source=ExtractionSource.PASSTHROUGH
+                if output.is_passthrough
+                else ExtractionSource.PROCESSED,
+            )
+            extraction.save_to_file()
+
+            return True
+        except Exception as e:
+            logger.error(
+                f"Error running extraction job for dataset item {job.doc.id}: {e}"
+            )
+            return False

kiln_ai/adapters/extractors/litellm_extractor.py

@@ -0,0 +1,386 @@
+import asyncio
+import hashlib
+import logging
+from pathlib import Path
+from typing import Any, List
+
+import litellm
+from litellm.types.utils import Choices, ModelResponse
+
+from kiln_ai.adapters.extractors.base_extractor import (
+    BaseExtractor,
+    ExtractionInput,
+    ExtractionOutput,
+)
+from kiln_ai.adapters.extractors.encoding import to_base64_url
+from kiln_ai.adapters.ml_model_list import built_in_models_from_provider
+from kiln_ai.adapters.provider_tools import LiteLlmCoreConfig
+from kiln_ai.datamodel.datamodel_enums import ModelProviderName
+from kiln_ai.datamodel.extraction import ExtractorConfig, ExtractorType, Kind
+from kiln_ai.utils.filesystem_cache import FilesystemCache
+from kiln_ai.utils.litellm import get_litellm_provider_info
+from kiln_ai.utils.pdf_utils import split_pdf_into_pages
+
+
+def max_pdf_page_concurrency_for_model(model_name: str) -> int:
+    # we assume each batch takes ~5s to complete (likely more in practice)
+    # lowest rate limit is 150 RPM for Tier 1 accounts for gemini-2.5-pro
+    if model_name == "gemini/gemini-2.5-pro":
+        return 2
+    # other models support at least 500 RPM for lowest tier accounts
+    return 5
+
+
+logger = logging.getLogger(__name__)
+
+MIME_TYPES_SUPPORTED = {
+    Kind.DOCUMENT: [
+        "application/pdf",
+        "text/plain",
+        "text/markdown",  # not officially listed, but works
+        "text/html",
+        "text/md",
+        "text/csv",
+    ],
+    Kind.IMAGE: [
+        "image/png",
+        "image/jpeg",
+        "image/jpg",
+    ],
+    Kind.VIDEO: [
+        "video/mp4",
+        "video/mov",  # the correct type is video/quicktime, but Google lists it as video/mov
+        "video/quicktime",
+    ],
+    Kind.AUDIO: [
+        "audio/wav",
+        "audio/mpeg",  # this is the official MP3 mimetype, audio/mp3 is often used but not correct
+        "audio/ogg",
+    ],
+}
+
+
+def encode_file_litellm_format(path: Path, mime_type: str) -> dict[str, Any]:
+    # There are different formats that LiteLLM supports, the docs are scattered
+    # and incomplete:
+    # - https://docs.litellm.ai/docs/completion/document_understanding#base64
+    # - https://docs.litellm.ai/docs/completion/vision#explicitly-specify-image-type
+
+    # this is the most generic format that seems to work for all / most mime types
+    if mime_type in [
+        "application/pdf",
+        "text/csv",
+        "text/html",
+        "text/markdown",
+        "text/plain",
+    ] or any(mime_type.startswith(m) for m in ["video/", "audio/"]):
+        pdf_bytes = path.read_bytes()
+        return {
+            "type": "file",
+            "file": {
+                "file_data": to_base64_url(mime_type, pdf_bytes),
+            },
+        }
+
+    # image has its own format (but also appears to work with the file format)
+    if mime_type.startswith("image/"):
+        image_bytes = path.read_bytes()
+        return {
+            "type": "image_url",
+            "image_url": {
+                "url": to_base64_url(mime_type, image_bytes),
+            },
+        }
+
+    raise ValueError(f"Unsupported MIME type: {mime_type} for {path}")
+
+
+class LitellmExtractor(BaseExtractor):
+    def __init__(
+        self,
+        extractor_config: ExtractorConfig,
+        litellm_core_config: LiteLlmCoreConfig,
+        filesystem_cache: FilesystemCache | None = None,
+    ):
+        if extractor_config.extractor_type != ExtractorType.LITELLM:
+            raise ValueError(
+                f"LitellmExtractor must be initialized with a litellm extractor_type config. Got {extractor_config.extractor_type}"
+            )
+
+        prompt_document = extractor_config.prompt_document()
+        if prompt_document is None or prompt_document == "":
+            raise ValueError(
+                "properties.prompt_document is required for LitellmExtractor"
+            )
+        prompt_video = extractor_config.prompt_video()
+        if prompt_video is None or prompt_video == "":
+            raise ValueError("properties.prompt_video is required for LitellmExtractor")
+        prompt_audio = extractor_config.prompt_audio()
+        if prompt_audio is None or prompt_audio == "":
+            raise ValueError("properties.prompt_audio is required for LitellmExtractor")
+        prompt_image = extractor_config.prompt_image()
+        if prompt_image is None or prompt_image == "":
+            raise ValueError("properties.prompt_image is required for LitellmExtractor")
+
+        self.filesystem_cache = filesystem_cache
+
+        super().__init__(extractor_config)
+        self.prompt_for_kind = {
+            Kind.DOCUMENT: prompt_document,
+            Kind.VIDEO: prompt_video,
+            Kind.AUDIO: prompt_audio,
+            Kind.IMAGE: prompt_image,
+        }
+
+        self.litellm_core_config = litellm_core_config
+
+    def pdf_page_cache_key(self, pdf_path: Path, page_number: int) -> str:
+        """
+        Generate a cache key for a page of a PDF. The PDF path must be the full path to the PDF file,
+        not the path to the page - since page path is temporary and changes on each run.
+        """
+        if self.extractor_config.id is None:
+            raise ValueError("Extractor config ID is required for PDF page cache key")
+
+        raw_key = f"{pdf_path.resolve()}::{page_number}"
+        digest = hashlib.md5(raw_key.encode("utf-8")).hexdigest()
+        return f"{self.extractor_config.id}_{digest}"
+
+    async def get_page_content_from_cache(
+        self, pdf_path: Path, page_number: int
+    ) -> str | None:
+        if self.filesystem_cache is None:
+            return None
+
+        page_bytes = await self.filesystem_cache.get(
+            self.pdf_page_cache_key(pdf_path, page_number)
+        )
+
+        if page_bytes is not None:
+            logger.debug(f"Cache hit for page {page_number} of {pdf_path}")
+            try:
+                return page_bytes.decode("utf-8")
+            except UnicodeDecodeError:
+                logger.warning(
+                    "Cached bytes for page %s of %s are not valid UTF-8; treating as miss.",
+                    page_number,
+                    pdf_path,
+                    exc_info=True,
+                )
+
+        logger.debug(f"Cache miss for page {page_number} of {pdf_path}")
+        return None
+
+    async def _extract_single_pdf_page(
+        self, pdf_path: Path, page_path: Path, prompt: str, page_number: int
+    ) -> str:
+        try:
+            page_input = ExtractionInput(
+                path=str(page_path), mime_type="application/pdf"
+            )
+            completion_kwargs = self._build_completion_kwargs(prompt, page_input)
+            response = await litellm.acompletion(**completion_kwargs)
+        except Exception as e:
+            raise RuntimeError(
+                f"Error extracting page {page_number} in file {page_path}: {e}"
+            ) from e
+
+        if (
+            not isinstance(response, ModelResponse)
+            or not response.choices
+            or len(response.choices) == 0
+            or not isinstance(response.choices[0], Choices)
+        ):
+            raise RuntimeError(
+                f"Expected ModelResponse with Choices for page {page_number}, got {type(response)}."
+            )
+
+        if response.choices[0].message.content is None:
+            raise ValueError(
+                f"No text returned from LiteLLM when extracting page {page_number}"
+            )
+
+        content = response.choices[0].message.content
+        if not content:
+            raise ValueError(
+                f"No text returned from extraction model when extracting page {page_number} for {page_path}"
+            )
+
+        if self.filesystem_cache is not None:
+            # we don't want to fail the whole extraction just because cache write fails
+            # as that would block the whole flow
+            try:
+                logger.debug(f"Caching page {page_number} of {page_path} in cache")
+                await self.filesystem_cache.set(
+                    self.pdf_page_cache_key(pdf_path, page_number),
+                    content.encode("utf-8"),
+                )
+            except Exception:
+                logger.warning(
+                    "Failed to cache page %s of %s; continuing without cache.",
+                    page_number,
+                    page_path,
+                    exc_info=True,
+                )
+
+        return content
+
+    async def _extract_pdf_page_by_page(self, pdf_path: Path, prompt: str) -> str:
+        async with split_pdf_into_pages(pdf_path) as page_paths:
+            page_outcomes: List[str | Exception | None] = [None] * len(page_paths)
+
+            extract_page_jobs: list = []
+            page_indices_for_jobs: list = []  # Track which page index each job corresponds to
+
+            # we extract from each page individually and then combine the results
+            # this ensures the model stays focused on the current page and does not
+            # start summarizing the later pages
+            for i, page_path in enumerate(page_paths):
+                page_content = await self.get_page_content_from_cache(pdf_path, i)
+                if page_content is not None:
+                    page_outcomes[i] = page_content
+                    continue
+
+                extract_page_jobs.append(
+                    self._extract_single_pdf_page(pdf_path, page_path, prompt, i)
+                )
+                page_indices_for_jobs.append(i)
+
+                if (
+                    len(extract_page_jobs)
+                    >= max_pdf_page_concurrency_for_model(self.litellm_model_slug())
+                    or i == len(page_paths) - 1
+                ):
+                    extraction_results = await asyncio.gather(
+                        *extract_page_jobs, return_exceptions=True
+                    )
+
+                    for batch_i, extraction_result in enumerate(extraction_results):
+                        page_index = page_indices_for_jobs[batch_i]
+                        # we let it continue even if there is an error - the success results will be cached
+                        # and can be reused on the next run
+                        if isinstance(extraction_result, Exception):
+                            page_outcomes[page_index] = extraction_result
+                        elif isinstance(extraction_result, str):
+                            page_outcomes[page_index] = extraction_result
+                        else:
+                            raise ValueError(
+                                f"Unexpected type {type(extraction_result)} for page {page_index}"
+                            )
+                    extract_page_jobs.clear()
+                    page_indices_for_jobs.clear()
+
+        exceptions: list[tuple[int, Exception]] = [
+            (page_index, result)
+            for page_index, result in enumerate(page_outcomes)
+            if isinstance(result, Exception)
+        ]
+        if len(exceptions) > 0:
+            msg = f"Error extracting PDF {pdf_path}: "
+            for page_index, exception in exceptions:
+                msg += f"Page {page_index}: {exception}\n"
+            raise RuntimeError(msg)
+
+        return "\n\n".join(
+            [outcome for outcome in page_outcomes if isinstance(outcome, str)]
+        )
+
+    def _get_kind_from_mime_type(self, mime_type: str) -> Kind | None:
+        for kind, mime_types in MIME_TYPES_SUPPORTED.items():
+            if mime_type in mime_types:
+                return kind
+        return None
+
+    def _build_completion_kwargs(
+        self, prompt: str, extraction_input: ExtractionInput
+    ) -> dict[str, Any]:
+        completion_kwargs = {
+            "model": self.litellm_model_slug(),
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": prompt},
+                        encode_file_litellm_format(
+                            Path(extraction_input.path), extraction_input.mime_type
+                        ),
+                    ],
+                }
+            ],
+        }
+
+        if self.litellm_core_config.base_url:
+            completion_kwargs["base_url"] = self.litellm_core_config.base_url
+
+        if self.litellm_core_config.default_headers:
+            completion_kwargs["default_headers"] = (
+                self.litellm_core_config.default_headers
+            )
+
+        if self.litellm_core_config.additional_body_options:
+            completion_kwargs.update(self.litellm_core_config.additional_body_options)
+
+        return completion_kwargs
+
+    async def _extract(self, extraction_input: ExtractionInput) -> ExtractionOutput:
+        kind = self._get_kind_from_mime_type(extraction_input.mime_type)
+        if kind is None:
+            raise ValueError(
+                f"Unsupported MIME type: {extraction_input.mime_type} for {extraction_input.path}"
+            )
+
+        prompt = self.prompt_for_kind.get(kind)
+        if prompt is None:
+            raise ValueError(f"No prompt found for kind: {kind}")
+
+        # special handling for PDFs - process each page individually
+        if extraction_input.mime_type == "application/pdf":
+            content = await self._extract_pdf_page_by_page(
+                Path(extraction_input.path), prompt
+            )
+            return ExtractionOutput(
+                is_passthrough=False,
+                content=content,
+                content_format=self.extractor_config.output_format,
+            )
+
+        completion_kwargs = self._build_completion_kwargs(prompt, extraction_input)
+
+        response = await litellm.acompletion(**completion_kwargs)
+
+        if (
+            not isinstance(response, ModelResponse)
+            or not response.choices
+            or len(response.choices) == 0
+            or not isinstance(response.choices[0], Choices)
+        ):
+            raise RuntimeError(
+                f"Expected ModelResponse with Choices, got {type(response)}."
+            )
+
+        if response.choices[0].message.content is None:
+            raise ValueError("No text returned from LiteLLM when extracting document")
+
+        return ExtractionOutput(
+            is_passthrough=False,
+            content=response.choices[0].message.content,
+            content_format=self.extractor_config.output_format,
+        )
+
+    def litellm_model_slug(self) -> str:
+        kiln_model_provider = built_in_models_from_provider(
+            ModelProviderName(self.extractor_config.model_provider_name),
+            self.extractor_config.model_name,
+        )
+
+        if kiln_model_provider is None:
+            raise ValueError(
+                f"Model provider {self.extractor_config.model_provider_name} not found in the list of built-in models"
+            )
+
+        # need to translate into LiteLLM model slug
+        litellm_provider_name = get_litellm_provider_info(
+            kiln_model_provider,
+        )
+
+        return litellm_provider_name.litellm_model_id