kiln-ai 0.20.1__py3-none-any.whl → 0.22.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