kiln-ai 0.8.1__py3-none-any.whl → 0.12.0__py3-none-any.whl

kiln_ai/adapters/__init__.py

@@ -3,31 +3,31 @@
 
 Adapters are used to connect Kiln to external systems, or to add new functionality to Kiln.
 
-BaseAdapter is extensible, and used for adding adapters that provide AI functionality. There's currently a LangChain adapter which provides a bridge to LangChain.
+Model adapters are used to call AI models, like Ollama, OpenAI, etc.
 
 The ml_model_list submodule contains a list of models that can be used for machine learning tasks. More can easily be added, but we keep a list here of models that are known to work well with Kiln's structured data and tool calling systems.
 
 The prompt_builders submodule contains classes that build prompts for use with the AI agents.
 
 The repair submodule contains an adapter for the repair task.
+
+The parser submodule contains parsers for the output of the AI models.
 """
 
 from . import (
-    base_adapter,
     data_gen,
     fine_tune,
-    langchain_adapters,
     ml_model_list,
+    model_adapters,
     prompt_builders,
     repair,
 )
 
 __all__ = [
-    "base_adapter",
-    "langchain_adapters",
+    "model_adapters",
+    "data_gen",
+    "fine_tune",
     "ml_model_list",
     "prompt_builders",
     "repair",
-    "data_gen",
-    "fine_tune",
 ]

kiln_ai/adapters/adapter_registry.py

@@ -1,21 +1,92 @@
+from os import getenv
+
 from kiln_ai import datamodel
-from kiln_ai.adapters.base_adapter import BaseAdapter
-from kiln_ai.adapters.langchain_adapters import LangchainAdapter
-from kiln_ai.adapters.prompt_builders import BasePromptBuilder
+from kiln_ai.adapters.ml_model_list import ModelProviderName
+from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig, BaseAdapter
+from kiln_ai.adapters.model_adapters.langchain_adapters import LangchainAdapter
+from kiln_ai.adapters.model_adapters.openai_model_adapter import (
+    OpenAICompatibleAdapter,
+    OpenAICompatibleConfig,
+)
+from kiln_ai.adapters.provider_tools import core_provider, openai_compatible_config
+from kiln_ai.datamodel import PromptId
+from kiln_ai.utils.config import Config
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 
 def adapter_for_task(
     kiln_task: datamodel.Task,
-    model_name: str | None = None,
-    provider: str | None = None,
-    prompt_builder: BasePromptBuilder | None = None,
-    tags: list[str] | None = None,
+    model_name: str,
+    provider: ModelProviderName,
+    prompt_id: PromptId | None = None,
+    base_adapter_config: AdapterConfig | None = None,
 ) -> BaseAdapter:
-    # We use langchain for everything right now, but can add any others here
+    # Get the provider to run. For things like the fine-tune provider, we want to run the underlying provider
+    core_provider_name = core_provider(model_name, provider)
+
+    match core_provider_name:
+        case ModelProviderName.openrouter:
+            return OpenAICompatibleAdapter(
+                kiln_task=kiln_task,
+                config=OpenAICompatibleConfig(
+                    base_url=getenv("OPENROUTER_BASE_URL")
+                    or "https://openrouter.ai/api/v1",
+                    api_key=Config.shared().open_router_api_key,
+                    model_name=model_name,
+                    provider_name=provider,
+                    default_headers={
+                        "HTTP-Referer": "https://getkiln.ai/openrouter",
+                        "X-Title": "KilnAI",
+                    },
+                ),
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+            )
+        case ModelProviderName.openai:
+            return OpenAICompatibleAdapter(
+                kiln_task=kiln_task,
+                config=OpenAICompatibleConfig(
+                    api_key=Config.shared().open_ai_api_key,
+                    model_name=model_name,
+                    provider_name=provider,
+                ),
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+            )
+        case ModelProviderName.openai_compatible:
+            config = openai_compatible_config(model_name)
+            return OpenAICompatibleAdapter(
+                kiln_task=kiln_task,
+                config=config,
+                prompt_id=prompt_id,
+                base_adapter_config=base_adapter_config,
+            )
+        # Use LangchainAdapter for the rest
+        case ModelProviderName.groq:
+            pass
+        case ModelProviderName.amazon_bedrock:
+            pass
+        case ModelProviderName.ollama:
+            pass
+        case ModelProviderName.fireworks_ai:
+            pass
+        # These are virtual providers that should have mapped to an actual provider in core_provider
+        case ModelProviderName.kiln_fine_tune:
+            raise ValueError(
+                "Fine tune is not a supported core provider. It should map to an actual provider."
+            )
+        case ModelProviderName.kiln_custom_registry:
+            raise ValueError(
+                "Custom openai compatible provider is not a supported core provider. It should map to an actual provider."
+            )
+        case _:
+            raise_exhaustive_enum_error(core_provider_name)
+
+    # We use langchain for all others right now, but moving off it as we touch anything.
     return LangchainAdapter(
         kiln_task,
         model_name=model_name,
         provider=provider,
-        prompt_builder=prompt_builder,
-        tags=tags,
+        prompt_id=prompt_id,
+        base_adapter_config=base_adapter_config,
     )

kiln_ai/adapters/data_gen/data_gen_task.py

@@ -55,7 +55,7 @@ class DataGenCategoriesTaskInput(BaseModel):
             num_subtopics=num_subtopics,
             human_guidance=human_guidance,
             existing_topics=existing_topics,
-            system_prompt=prompt_builder.build_prompt(),
+            system_prompt=prompt_builder.build_prompt(include_json_instructions=False),
         )
 
 
@@ -132,7 +132,7 @@ class DataGenSampleTaskInput(BaseModel):
             topic=topic,
             num_samples=num_samples,
             human_guidance=human_guidance,
-            system_prompt=prompt_builder.build_prompt(),
+            system_prompt=prompt_builder.build_prompt(include_json_instructions=False),
         )
 
 
@@ -163,7 +163,7 @@ def list_json_schema_for_task(task: Task) -> str:
         "required": ["generated_samples"],
     }
 
-    return json.dumps(top_level_schema)
+    return json.dumps(top_level_schema, ensure_ascii=False)
 
 
 class DataGenSampleTask(Task, parent_of={}):
@@ -183,3 +183,21 @@ class DataGenSampleTask(Task, parent_of={}):
             input_json_schema=json.dumps(DataGenSampleTaskInput.model_json_schema()),
             output_json_schema=list_json_schema_for_task(target_task),
         )
+
+
+def wrap_task_with_guidance(original_instruction: str, guidance: str) -> str:
+    """Wrap the original instruction with human guidance.
+
+    Args:
+        original_instruction: The original instruction to wrap
+        guidance: The human guidance to wrap the instruction with
+    """
+    return f"""{original_instruction}
+
+# Special Instructions
+
+The above instructions are the original instructions for this task. For this execution, we've been given additional instructions. Follow both, but prioritize the additional instructions when they conflict. The additional instructions are:
+<additional_instructions>
+{guidance}
+</additional_instructions>
+"""

kiln_ai/adapters/data_gen/test_data_gen_task.py

@@ -180,7 +180,7 @@ def test_data_gen_sample_task_initialization(base_task):
     }
 
 
-def test_list_json_schema_for_task_with_output_schema(base_task):
+def test_list_json_schema_for_task_with_input_schema(base_task):
     # Arrange
     base_task.input_json_schema = json.dumps(
         {
@@ -202,9 +202,29 @@ def test_list_json_schema_for_task_with_output_schema(base_task):
     assert generated_samples_schema["items"]["properties"]["age"]["type"] == "integer"
 
 
-def test_list_json_schema_for_task_without_output_schema(base_task):
+def test_list_json_schema_for_task_with_input_schema_non_ascii(base_task):
     # Arrange
-    base_task.output_json_schema = None
+    base_task.input_json_schema = json.dumps(
+        {
+            "type": "object",
+            "properties": {
+                "名字": {"type": "string"},
+                "年齢": {"type": "integer"},
+            },
+        }
+    )
+
+    # Act
+    schema = list_json_schema_for_task(base_task)
+
+    # Assert
+    assert "名字" in schema
+    assert "年齢" in schema
+
+
+def test_list_json_schema_for_task_without_input_schema(base_task):
+    # Arrange
+    base_task.input_json_schema = None
 
     # Act
     schema = list_json_schema_for_task(base_task)

kiln_ai/adapters/eval/base_eval.py

@@ -0,0 +1,164 @@
+import json
+from abc import abstractmethod
+from typing import Dict
+
+from kiln_ai.adapters.adapter_registry import adapter_for_task
+from kiln_ai.adapters.ml_model_list import ModelProviderName
+from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig
+from kiln_ai.datamodel.eval import Eval, EvalConfig, EvalScores
+from kiln_ai.datamodel.json_schema import validate_schema
+from kiln_ai.datamodel.task import RunConfig, TaskOutputRatingType, TaskRun
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+
+class BaseEval:
+    """
+    Base class for all evals/evaluators.
+
+    Should be subclassed, and the run_eval method implemented.
+    """
+
+    def __init__(self, eval_config: EvalConfig, run_config: RunConfig | None):
+        self.eval_config = eval_config
+        eval = eval_config.parent_eval()
+        if not eval:
+            raise ValueError("Eval config must have a parent eval")
+        self.eval = eval
+        task = self.eval.parent_task()
+        if not task:
+            raise ValueError("Eval must have a parent task")
+        self.target_task = task
+        self.score_schema = BaseEval.build_score_schema(eval, allow_float_scores=True)
+        self.run_config = run_config
+
+    def model_and_provider(self) -> tuple[str, ModelProviderName]:
+        model_name = self.eval_config.model_name
+        provider = self.eval_config.model_provider
+        if (
+            not model_name
+            or not provider
+            or not isinstance(model_name, str)
+            or not isinstance(provider, str)
+            or provider not in ModelProviderName.__members__
+        ):
+            raise ValueError(
+                "Model name and provider must be set in the eval config model properties"
+            )
+
+        return model_name, ModelProviderName(provider)
+
+    async def run_task_and_eval(
+        self, input: str
+    ) -> tuple[TaskRun, EvalScores, Dict[str, str] | None]:
+        """
+        Runs the task on the provided run_config to generate fresh output, then runs the eval on that output.
+        """
+        if self.run_config is None:
+            raise ValueError("Run config is required for run_task_and_eval")
+
+        run_adapter = adapter_for_task(
+            self.target_task,
+            self.run_config.model_name,
+            ModelProviderName(self.run_config.model_provider_name),
+            base_adapter_config=AdapterConfig(allow_saving=False),
+        )
+
+        # Parse structured input if needed
+        parsed_input = input
+        if self.target_task.output_json_schema is not None:
+            parsed_input = json.loads(input)
+
+        # we don't save by default here. We'll save manually after validating the output
+        run_output = await run_adapter.invoke(parsed_input)
+
+        eval_output, intermediate_outputs = await self.run_eval(run_output)
+        validate_schema(eval_output, self.score_schema)
+
+        return run_output, eval_output, intermediate_outputs
+
+    @abstractmethod
+    async def run_eval(
+        self, task_run: TaskRun
+    ) -> tuple[EvalScores, Dict[str, str] | None]:
+        """
+        Runs the eval on the given task run.
+
+        Returns a dictionary of scores which should conform to the score schema, and a dictionary of intermediate outputs (eval thinking).
+        """
+        pass
+
+    @classmethod
+    def build_score_schema(cls, eval: Eval, allow_float_scores: bool = False) -> str:
+        """
+        Build a JSON schema for the scoring output of the task requirements
+
+        We allow 2 modes: allow_float_scores=True and allow_float_scores=False.
+
+        allow_float_scores=False is used for the call to the model, and forces the model into selecting into discrete rating options (int 1-5, pass-fail, etc).
+        allow_float_scores=True is used for final score output (for example, after we take a g-eval weighting of the model's logprobs). A pass/fail rating might return 0.75 for likely pass (as opposed to 0.99 for near certain pass), or a 1-5 score might return 3.75.
+        """
+
+        # Note: python maintains order, which is good as we want the user defined order, and overall last
+        properties = {}
+        for output_score in eval.output_scores:
+            output_score_json_key = output_score.json_key()
+
+            if len(output_score_json_key) == 0:
+                raise ValueError(
+                    f"Invalid output score name: {output_score.name}. Can not be used as JSON schema key."
+                )
+            property: dict[str, str | int | float | list[str] | list[int]] = {
+                "title": output_score.name,
+            }
+            match output_score.type:
+                case TaskOutputRatingType.five_star:
+                    if allow_float_scores:
+                        property["type"] = "number"
+                        property["minimum"] = 1
+                        property["maximum"] = 5
+                    else:
+                        property["enum"] = [1, 2, 3, 4, 5]
+
+                    property["description"] = (
+                        f"{output_score.instruction}\n\nThe rating should be between 1 and 5, with 1 being the worst and 5 being the best."
+                    )
+                case TaskOutputRatingType.pass_fail:
+                    if allow_float_scores:
+                        property["type"] = "number"
+                        property["minimum"] = 0
+                        property["maximum"] = 1
+                        property["description"] = (
+                            f"{output_score.instruction}\n\nThe rating should be between 0 and 1, with 0 being a failure and 1 being a pass."
+                        )
+                    else:
+                        property["enum"] = ["pass", "fail"]
+                        property["description"] = (
+                            f"{output_score.instruction}\n\nThe rating should be either 'pass' or 'fail'."
+                        )
+                case TaskOutputRatingType.pass_fail_critical:
+                    if allow_float_scores:
+                        property["type"] = "number"
+                        property["minimum"] = -1
+                        property["maximum"] = 1
+                        property["description"] = (
+                            f"{output_score.instruction}\n\nThe rating should be between -1 and 1, with 1 being a pass, 0 being a failure, and -1 being a critical failure (very severe failure)."
+                        )
+                    else:
+                        property["enum"] = ["pass", "fail", "critical"]
+                        property["description"] = (
+                            f"{output_score.instruction}\n\nThe rating should be either 'pass', 'fail', or 'critical' where critical a very severe failure."
+                        )
+                case TaskOutputRatingType.custom:
+                    # Skip custom rating types in evals
+                    continue
+                case _:
+                    raise_exhaustive_enum_error(output_score.type)
+
+            properties[output_score_json_key] = property
+
+        schema = {
+            "type": "object",
+            "properties": properties,
+            "required": list(properties.keys()),
+        }
+        return json.dumps(schema, ensure_ascii=False)

kiln_ai/adapters/eval/eval_runner.py

@@ -0,0 +1,267 @@
+import asyncio
+import logging
+from dataclasses import dataclass
+from typing import AsyncGenerator, Dict, List, Literal, Set
+
+from kiln_ai.adapters.eval.base_eval import BaseEval
+from kiln_ai.adapters.eval.registry import eval_adapter_from_type
+from kiln_ai.datamodel.basemodel import ID_TYPE
+from kiln_ai.datamodel.dataset_filters import dataset_filter_from_id
+from kiln_ai.datamodel.eval import EvalConfig, EvalRun, EvalScores
+from kiln_ai.datamodel.task import TaskRunConfig
+from kiln_ai.datamodel.task_run import TaskRun
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class EvalJob:
+    item: TaskRun
+    type: Literal["task_run_eval", "eval_config_eval"]
+    # If type == "task_run_eval", both of these should be set. If type == "eval_config_eval", only eval_config should be set.
+    eval_config: EvalConfig
+    task_run_config: TaskRunConfig | None = None
+
+
+@dataclass
+class EvalProgress:
+    complete: int | None = None
+    total: int | None = None
+    errors: int | None = None
+
+
+class EvalRunner:
+    """
+    Runs an eval. Async execution is supported to make it faster when using remote/fast model providers.
+
+    Can run an eval in 2 modes:
+    1) eval_config_eval: evaluate an eval config using existing dataset items.
+    2) task_run_eval: evaluate a range of task run configs, generating new run output using existing dataset item input.
+    """
+
+    def __init__(
+        self,
+        eval_configs: List[EvalConfig],
+        run_configs: List[TaskRunConfig] | None,
+        eval_run_type: Literal["eval_config_eval", "task_run_eval"],
+    ):
+        if len(eval_configs) == 0:
+            raise ValueError("Eval runner requires at least one eval config")
+        target_eval = eval_configs[0].parent_eval()
+        if target_eval is None:
+            raise ValueError("Eval config requires a parent eval")
+        for eval_config in eval_configs:
+            parent_eval = eval_config.parent_eval()
+            if parent_eval is None:
+                raise ValueError("Eval config requires a parent eval")
+            if parent_eval.id != target_eval.id:
+                raise ValueError("All eval configs must have the same parent eval")
+
+        target_task = target_eval.parent_task()
+        if target_task is None:
+            raise ValueError("Eval config requires a (grand)parent task")
+
+        # Check that run_configs is compatible
+        if eval_run_type == "task_run_eval":
+            if run_configs is None or len(run_configs) == 0:
+                raise ValueError("Task run eval requires run configs")
+            for run_config in run_configs:
+                parent_task = run_config.parent_task()
+                if parent_task is None:
+                    raise ValueError("All run configs must have a parent task")
+                if parent_task.id != target_task.id:
+                    raise ValueError(
+                        "Run config is not for the same task as the eval configs"
+                    )
+        else:
+            if run_configs is not None:
+                raise ValueError("Mode 'eval_config_eval' does not support run configs")
+
+        self.eval_run_type = eval_run_type
+        self.eval_configs = eval_configs
+        self.run_configs = run_configs
+        self.task = target_task
+        self.eval = target_eval
+
+    def collect_tasks(self) -> List[EvalJob]:
+        if self.eval_run_type == "eval_config_eval":
+            return self.collect_tasks_for_eval_config_eval()
+        else:
+            return self.collect_tasks_for_task_run_eval()
+
+    def collect_tasks_for_eval_config_eval(self) -> List[EvalJob]:
+        """
+        Collect all jobs for this run, excluding any that have already been run.
+
+        This variant is used for mode "eval_config_eval", using existing dataset run data (input/output).
+
+        The tasks:
+        - should be in the eval config set filter
+        - should not have already been run for this eval config + dataset item pair
+        """
+        filter = dataset_filter_from_id(self.eval.eval_configs_filter_id)
+
+        # already_run[eval_config_id][dataset_id]
+        already_run: Dict[ID_TYPE, Set[ID_TYPE]] = {}
+        for eval_config in self.eval_configs:
+            already_run[eval_config.id] = set()
+            for run in eval_config.runs(readonly=True):
+                already_run[eval_config.id].add(run.dataset_id)
+
+        return [
+            EvalJob(
+                item=task_run,
+                eval_config=eval_config,
+                type="eval_config_eval",
+            )
+            for task_run in self.task.runs(readonly=True)
+            if filter(task_run)
+            for eval_config in self.eval_configs
+            if task_run.id not in already_run[eval_config.id]
+        ]
+
+    def collect_tasks_for_task_run_eval(self) -> List[EvalJob]:
+        """
+        Collect all jobs for this run, excluding any that have already been run.
+
+        This variant is used for mode "task_run_eval", generating new run output using existing dataset item input.
+
+        The tasks:
+        - should be in the eval set filter
+        - should not have already been run for this eval config + run config + dataset item
+        """
+        filter = dataset_filter_from_id(self.eval.eval_set_filter_id)
+
+        # already_run[eval_config_id][run_config_id][dataset_id]
+        already_run: Dict[ID_TYPE, Dict[ID_TYPE, Set[ID_TYPE]]] = {}
+        for eval_config in self.eval_configs:
+            already_run[eval_config.id] = {}
+            for run_config in self.run_configs or []:
+                already_run[eval_config.id][run_config.id] = set()
+            for run in eval_config.runs(readonly=True):
+                if run.task_run_config_id is not None:
+                    already_run[eval_config.id][run.task_run_config_id].add(
+                        run.dataset_id
+                    )
+
+        return [
+            EvalJob(
+                item=task_run,
+                task_run_config=run_config,
+                type="task_run_eval",
+                eval_config=eval_config,
+            )
+            for task_run in self.task.runs(readonly=True)
+            if filter(task_run)
+            for eval_config in self.eval_configs
+            for run_config in self.run_configs or []
+            if task_run.id not in already_run[eval_config.id][run_config.id]
+        ]
+
+    async def run(self, concurrency: int = 25) -> AsyncGenerator[EvalProgress, None]:
+        """
+        Runs the configured eval run with parallel workers and yields progress updates.
+        """
+        jobs = self.collect_tasks()
+
+        complete = 0
+        errors = 0
+        total = len(jobs)
+
+        # Send initial status
+        yield EvalProgress(complete=complete, total=total, errors=errors)
+
+        worker_queue: asyncio.Queue[EvalJob] = asyncio.Queue()
+        for job in jobs:
+            worker_queue.put_nowait(job)
+
+        # simple status queue to return progress. True=success, False=error
+        status_queue: asyncio.Queue[bool] = asyncio.Queue()
+
+        workers = []
+        for i in range(concurrency):
+            task = asyncio.create_task(self.run_worker(worker_queue, status_queue))
+            workers.append(task)
+
+        # Send status updates until workers are done, and they are all sent
+        while not status_queue.empty() or not all(worker.done() for worker in workers):
+            try:
+                # Use timeout to prevent hanging if all workers complete
+                # between our while condition check and get()
+                success = await asyncio.wait_for(status_queue.get(), timeout=0.1)
+                if success:
+                    complete += 1
+                else:
+                    errors += 1
+
+                yield EvalProgress(complete=complete, total=total, errors=errors)
+            except asyncio.TimeoutError:
+                # Timeout is expected, just continue to recheck worker status
+                # Don't love this but beats sentinels for reliability
+                continue
+
+        # These are redundant, but keeping them will catch async errors
+        await asyncio.gather(*workers)
+        await worker_queue.join()
+
+    async def run_worker(
+        self, worker_queue: asyncio.Queue[EvalJob], status_queue: asyncio.Queue[bool]
+    ):
+        while True:
+            try:
+                job = worker_queue.get_nowait()
+            except asyncio.QueueEmpty:
+                # worker can end when the queue is empty
+                break
+            try:
+                success = await self.run_job(job)
+                await status_queue.put(success)
+            finally:
+                # Always mark the dequeued task as done, even on exceptions
+                worker_queue.task_done()
+
+    async def run_job(self, job: EvalJob) -> bool:
+        try:
+            # Create the evaluator for this eval config/run config pair
+            evaluator = eval_adapter_from_type(job.eval_config.config_type)(
+                job.eval_config,
+                job.task_run_config.run_config() if job.task_run_config else None,
+            )
+            if not isinstance(evaluator, BaseEval):
+                raise ValueError("Not able to create evaluator from eval config")
+
+            task_output: str | None = None
+            scores: EvalScores | None = None
+            intermediate_outputs: Dict[str, str] | None = None
+            if job.type == "eval_config_eval":
+                # Eval config eval, we use the saved input from the task run, not invoking the task again
+                scores, intermediate_outputs = await evaluator.run_eval(job.item)
+                task_output = job.item.output.output
+            else:
+                # Task run eval, we invoke the task again to get a fresh output
+                (
+                    result_task_run,
+                    scores,
+                    intermediate_outputs,
+                ) = await evaluator.run_task_and_eval(job.item.input)
+                task_output = result_task_run.output.output
+
+            # Save the job result
+            eval_run = EvalRun(
+                parent=job.eval_config,
+                task_run_config_id=job.task_run_config.id
+                if job.task_run_config
+                else None,
+                dataset_id=job.item.id,
+                eval_config_eval=job.type == "eval_config_eval",
+                scores=scores,
+                input=job.item.input,
+                output=task_output,
+                intermediate_outputs=intermediate_outputs,
+            )
+            eval_run.save_to_file()
+
+            return True
+        except Exception as e:
+            logger.error(f"Error running eval job for dataset item {job.item.id}: {e}")
+            return False