kiln-ai 0.16.0__py3-none-any.whl → 0.18.0__py3-none-any.whl

kiln_ai/adapters/data_gen/test_data_gen_task.py

@@ -3,6 +3,10 @@ import json
 import pytest
 
 from kiln_ai.adapters.adapter_registry import adapter_for_task
+from kiln_ai.adapters.data_gen.data_gen_prompts import (
+    generate_sample_generation_prompt,
+    generate_topic_tree_prompt,
+)
 from kiln_ai.adapters.data_gen.data_gen_task import (
     DataGenCategoriesTask,
     DataGenCategoriesTaskInput,
@@ -14,6 +18,7 @@ from kiln_ai.adapters.data_gen.data_gen_task import (
 from kiln_ai.adapters.provider_tools import get_model_and_provider
 from kiln_ai.adapters.test_prompt_adaptors import get_all_models_and_providers
 from kiln_ai.datamodel import Project, Task
+from kiln_ai.datamodel.task import RunConfigProperties
 
 
 @pytest.fixture
@@ -32,22 +37,19 @@ def test_data_gen_categories_task_input_initialization(base_task):
     # Arrange
     node_path = ["root", "branch", "leaf"]
     num_subtopics = 4
-    human_guidance = "Test guidance"
 
     # Act
     input_model = DataGenCategoriesTaskInput.from_task(
         task=base_task,
         node_path=node_path,
         num_subtopics=num_subtopics,
-        human_guidance=human_guidance,
     )
 
     # Assert
-    assert input_model.node_path == node_path
-    assert input_model.num_subtopics == num_subtopics
-    assert input_model.human_guidance == human_guidance
-    assert isinstance(input_model.system_prompt, str)
-    assert "Reply like a cowboy" in input_model.system_prompt
+    assert input_model.kiln_data_gen_topic_path == node_path
+    assert input_model.kiln_data_gen_num_subtopics == num_subtopics
+    assert isinstance(input_model.kiln_data_gen_system_prompt, str)
+    assert "Reply like a cowboy" in input_model.kiln_data_gen_system_prompt
 
 
 def test_data_gen_categories_task_input_default_values(base_task):
@@ -55,14 +57,13 @@ def test_data_gen_categories_task_input_default_values(base_task):
     input_model = DataGenCategoriesTaskInput.from_task(task=base_task)
 
     # Assert
-    assert input_model.num_subtopics == 6
-    assert input_model.human_guidance is None
-    assert input_model.node_path == []
+    assert input_model.kiln_data_gen_num_subtopics == 6
+    assert input_model.kiln_data_gen_topic_path == []
 
 
 def test_data_gen_categories_task_initialization():
     # Act
-    task = DataGenCategoriesTask()
+    task = DataGenCategoriesTask(gen_type="training", guidance="Test guidance")
 
     # Assert
     assert task.name == "DataGen"
@@ -71,11 +72,16 @@ def test_data_gen_categories_task_initialization():
     assert task.instruction is not None
     assert isinstance(task.input_json_schema, str)
     assert isinstance(task.output_json_schema, str)
+    assert "I want to train a large language model" in task.instruction
+    assert "Test guidance" in task.instruction
 
 
 def test_data_gen_categories_task_schemas():
     # Act
-    task = DataGenCategoriesTask()
+    task = DataGenCategoriesTask(gen_type="eval", guidance="Test guidance")
+
+    assert "I want to evaluate a large language model" in task.instruction
+    assert "Test guidance" in task.instruction
 
     # Assert
     input_schema = json.loads(task.input_json_schema)
@@ -85,12 +91,14 @@ def test_data_gen_categories_task_schemas():
     assert isinstance(output_schema, dict)
     assert output_schema["type"] == "object"
     assert output_schema["properties"]["subtopics"]["type"] == "array"
-    assert input_schema["properties"]["node_path"]["type"] == "array"
-    assert input_schema["properties"]["num_subtopics"]["type"] == "integer"
+    assert input_schema["properties"]["kiln_data_gen_topic_path"]["type"] == "array"
+    assert (
+        input_schema["properties"]["kiln_data_gen_num_subtopics"]["type"] == "integer"
+    )
     assert set(input_schema["required"]) == {
-        "node_path",
-        "num_subtopics",
-        "system_prompt",
+        "kiln_data_gen_topic_path",
+        "kiln_data_gen_num_subtopics",
+        "kiln_data_gen_system_prompt",
     }
 
 
@@ -105,13 +113,17 @@ async def test_data_gen_all_models_providers(
         # pass if the model doesn't support data gen (testing the support flag is part of this)
         return
 
-    data_gen_task = DataGenCategoriesTask()
+    data_gen_task = DataGenCategoriesTask(gen_type="training", guidance=None)
     data_gen_input = DataGenCategoriesTaskInput.from_task(base_task, num_subtopics=6)
 
     adapter = adapter_for_task(
         data_gen_task,
-        model_name=model_name,
-        provider=provider_name,
+        run_config_properties=RunConfigProperties(
+            model_name=model_name,
+            model_provider_name=provider_name,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="unknown",
+        ),
     )
 
     input_dict = data_gen_input.model_dump()
@@ -126,22 +138,19 @@ def test_data_gen_sample_task_input_initialization(base_task):
     # Arrange
     topic = ["cowboys", "hats"]
     num_samples = 4
-    human_guidance = "Test guidance"
 
     # Act
     input_model = DataGenSampleTaskInput.from_task(
         task=base_task,
         topic=topic,
         num_samples=num_samples,
-        human_guidance=human_guidance,
     )
 
     # Assert
-    assert input_model.topic == topic
-    assert input_model.num_samples == num_samples
-    assert input_model.human_guidance == human_guidance
-    assert isinstance(input_model.system_prompt, str)
-    assert "Reply like a cowboy" in input_model.system_prompt
+    assert input_model.kiln_data_gen_topic_path == topic
+    assert input_model.kiln_data_gen_num_samples == num_samples
+    assert isinstance(input_model.kiln_data_gen_system_prompt, str)
+    assert "Reply like a cowboy" in input_model.kiln_data_gen_system_prompt
 
 
 def test_data_gen_sample_task_input_default_values(base_task):
@@ -149,20 +158,23 @@ def test_data_gen_sample_task_input_default_values(base_task):
     input_model = DataGenSampleTaskInput.from_task(task=base_task)
 
     # Assert
-    assert input_model.num_samples == 8
-    assert input_model.human_guidance is None
-    assert input_model.topic == []
+    assert input_model.kiln_data_gen_num_samples == 8
+    assert input_model.kiln_data_gen_topic_path == []
 
 
 def test_data_gen_sample_task_initialization(base_task):
     # Act
-    task = DataGenSampleTask(target_task=base_task)
+    task = DataGenSampleTask(
+        target_task=base_task, gen_type="eval", guidance="Test guidance"
+    )
 
     # Assert
     assert task.name == "DataGenSample"
     assert isinstance(task.parent, Project)
     assert task.description is not None
     assert task.instruction is not None
+    assert "I want to evaluate a large language model" in task.instruction
+    assert "Test guidance" in task.instruction
 
     input_schema = json.loads(task.input_json_schema)
     output_schema = json.loads(task.output_json_schema)
@@ -171,12 +183,12 @@ def test_data_gen_sample_task_initialization(base_task):
     assert isinstance(output_schema, dict)
     assert output_schema["type"] == "object"
     assert output_schema["properties"]["generated_samples"]["type"] == "array"
-    assert input_schema["properties"]["topic"]["type"] == "array"
-    assert input_schema["properties"]["num_samples"]["type"] == "integer"
+    assert input_schema["properties"]["kiln_data_gen_topic_path"]["type"] == "array"
+    assert input_schema["properties"]["kiln_data_gen_num_samples"]["type"] == "integer"
     assert set(input_schema["required"]) == {
-        "topic",
-        "num_samples",
-        "system_prompt",
+        "kiln_data_gen_topic_path",
+        "kiln_data_gen_num_samples",
+        "kiln_data_gen_system_prompt",
     }
 
 
@@ -254,8 +266,12 @@ async def test_data_gen_sample_all_models_providers(
 
     adapter = adapter_for_task(
         data_gen_task,
-        model_name=model_name,
-        provider=provider_name,
+        run_config_properties=RunConfigProperties(
+            model_name=model_name,
+            model_provider_name=provider_name,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="unknown",
+        ),
     )
 
     input_dict = data_gen_input.model_dump()
@@ -304,8 +320,12 @@ async def test_data_gen_sample_all_models_providers_with_structured_output(
 
     adapter = adapter_for_task(
         data_gen_task,
-        model_name=model_name,
-        provider=provider_name,
+        run_config_properties=RunConfigProperties(
+            model_name=model_name,
+            model_provider_name=provider_name,
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="unknown",
+        ),
     )
 
     input_dict = data_gen_input.model_dump()
@@ -319,3 +339,273 @@ async def test_data_gen_sample_all_models_providers_with_structured_output(
         assert "tweet" in sample
         assert isinstance(sample["username"], str)
         assert isinstance(sample["tweet"], str)
+
+
+def test_generate_topic_tree_prompt_training_type():
+    """Test generate_topic_tree_prompt with gen_type='training'"""
+    # Act
+    prompt = generate_topic_tree_prompt(gen_type="training")
+
+    # Assert
+    assert isinstance(prompt, str)
+    assert (
+        "I want to train a large language model and you should help me generate training data for it."
+        in prompt
+    )
+    assert "## Task Description" in prompt
+    assert "Your job is the following:" in prompt
+    assert "## Next Step" in prompt
+    assert "When generating subtopics, remain somewhat vague." in prompt
+    assert "The guidance is:" not in prompt  # Should not have specific guidance
+
+
+def test_generate_topic_tree_prompt_eval_type():
+    """Test generate_topic_tree_prompt with gen_type='eval'"""
+    # Act
+    prompt = generate_topic_tree_prompt(gen_type="eval")
+
+    # Assert
+    assert isinstance(prompt, str)
+    assert (
+        "I want to evaluate a large language model and you should help me generate eval data for it."
+        in prompt
+    )
+    assert "## Task Description" in prompt
+    assert "Your job is the following:" in prompt
+    assert "## Next Step" in prompt
+    assert "When generating subtopics, remain somewhat vague." in prompt
+    assert "The guidance is:" not in prompt  # Should not have specific guidance
+
+
+def test_generate_topic_tree_prompt_with_guidance():
+    """Test generate_topic_tree_prompt with guidance provided"""
+    # Arrange
+    guidance = "Focus on technical topics related to artificial intelligence and machine learning"
+
+    # Act
+    prompt = generate_topic_tree_prompt(gen_type="training", guidance=guidance)
+
+    # Assert
+    assert isinstance(prompt, str)
+    assert (
+        "I want to train a large language model and you should help me generate training data for it."
+        in prompt
+    )
+    assert "## Custom Guidance" in prompt
+    assert f"<guidance>\n{guidance}\n</guidance>" in prompt
+    assert (
+        "When generating subtopics, remain somewhat vague." not in prompt
+    )  # Should not have default guidance
+
+
+def test_generate_topic_tree_prompt_with_empty_guidance():
+    """Test generate_topic_tree_prompt with empty string guidance"""
+    # Act
+    prompt = generate_topic_tree_prompt(gen_type="eval", guidance="")
+
+    # Assert
+    assert isinstance(prompt, str)
+    assert (
+        "I want to evaluate a large language model and you should help me generate eval data for it."
+        in prompt
+    )
+    assert "## Specific Guidance" not in prompt
+    assert (
+        "When generating subtopics, remain somewhat vague." in prompt
+    )  # Should have default guidance
+
+
+def test_generate_topic_tree_prompt_contains_examples():
+    """Test that the prompt contains the expected examples"""
+    # Act
+    prompt = generate_topic_tree_prompt(gen_type="training")
+
+    # Assert
+    # Check for news examples
+    assert "News Topics" in prompt
+    assert "Sports" in prompt
+    assert "Football" in prompt
+    assert "College Football" in prompt
+    assert "Entertainment" in prompt
+    assert "Tom Hanks" in prompt
+
+    # Check for smalltalk examples
+    assert "Small Talk Topics" in prompt
+    assert "Weather" in prompt
+    assert "Family" in prompt
+    assert "Hobbies" in prompt
+    assert "Cooking" in prompt
+    assert "Asian Food" in prompt
+
+
+def test_generate_topic_tree_prompt_contains_required_sections():
+    """Test that the prompt contains all required sections"""
+    # Act
+    prompt = generate_topic_tree_prompt(gen_type="training")
+
+    # Assert
+    assert "## Task Description" in prompt
+    assert "## Next Step" in prompt
+    assert "system_prompt" in prompt
+    assert "kiln_data_gen_topic_path" in prompt
+    assert "kiln_data_gen_num_subtopics" in prompt
+    assert "existing_topics" in prompt
+
+
+def test_generate_topic_tree_prompt_structure_consistency():
+    """Test that the prompt structure is consistent between training and eval types"""
+    # Act
+    training_prompt = generate_topic_tree_prompt(gen_type="training")
+    eval_prompt = generate_topic_tree_prompt(gen_type="eval")
+
+    # Assert
+    # Both should have the same structure, just different goal descriptions
+    assert "## Task Description" in training_prompt
+    assert "## Task Description" in eval_prompt
+    assert "## Next Step" in training_prompt
+    assert "## Next Step" in eval_prompt
+
+    # The main difference should be in the goal description
+    assert "train a large language model" in training_prompt
+    assert "evaluate a large language model" in eval_prompt
+    assert "generate training data" in training_prompt
+    assert "generate eval data" in eval_prompt
+
+
+def test_generate_sample_generation_prompt_training_type():
+    """Test generate_sample_generation_prompt with gen_type='training'"""
+    # Act
+    prompt = generate_sample_generation_prompt(gen_type="training")
+
+    # Assert
+    assert isinstance(prompt, str)
+    assert (
+        "I want to train a large language model and you should help me generate training data for it."
+        in prompt
+    )
+    assert "## Task Description" in prompt
+    assert "Your job is to generate a list of potential inputs" in prompt
+    assert "The guidance is:" not in prompt  # Should not have specific guidance
+
+
+def test_generate_sample_generation_prompt_eval_type():
+    """Test generate_sample_generation_prompt with gen_type='eval'"""
+    # Act
+    prompt = generate_sample_generation_prompt(gen_type="eval")
+
+    # Assert
+    assert isinstance(prompt, str)
+    assert (
+        "I want to evaluate a large language model and you should help me generate eval data for it."
+        in prompt
+    )
+    assert "## Task Description" in prompt
+    assert "Your job is to generate a list of potential inputs" in prompt
+    assert "The guidance is:" not in prompt  # Should not have specific guidance
+
+
+def test_generate_sample_generation_prompt_with_guidance():
+    """Test generate_sample_generation_prompt with guidance provided"""
+    # Arrange
+    guidance = "Focus on generating diverse examples with varying complexity levels"
+
+    # Act
+    prompt = generate_sample_generation_prompt(gen_type="training", guidance=guidance)
+
+    # Assert
+    assert isinstance(prompt, str)
+    assert (
+        "I want to train a large language model and you should help me generate training data for it."
+        in prompt
+    )
+    assert "## Custom Guidance" in prompt
+    assert f"<guidance>\n{guidance}\n</guidance>" in prompt
+
+
+def test_generate_sample_generation_prompt_with_empty_guidance():
+    """Test generate_sample_generation_prompt with empty string guidance"""
+    # Act
+    prompt = generate_sample_generation_prompt(gen_type="eval", guidance="")
+
+    # Assert
+    assert isinstance(prompt, str)
+    assert (
+        "I want to evaluate a large language model and you should help me generate eval data for it."
+        in prompt
+    )
+    assert "## Specific Guidance" not in prompt
+
+
+def test_generate_sample_generation_prompt_contains_examples():
+    """Test that the prompt contains the expected examples"""
+    # Act
+    prompt = generate_sample_generation_prompt(gen_type="training")
+
+    # Assert
+    # Check for the tweet classification example
+    assert "You are an assistant that classifies the tone of a tweet" in prompt
+    assert "positive" in prompt
+    assert "negative" in prompt
+    assert "neutral" in prompt
+    assert "Technology" in prompt
+    assert "New iPhone Event" in prompt
+    assert "New iPhone looks amazing! I need that camera." in prompt
+    assert "Another boring event from Apple." in prompt
+
+
+def test_generate_sample_generation_prompt_contains_required_sections():
+    """Test that the prompt contains all required sections"""
+    # Act
+    prompt = generate_sample_generation_prompt(gen_type="training")
+
+    # Assert
+    assert "## Task Description" in prompt
+    assert "system_prompt" in prompt
+    assert "topic" in prompt
+    assert "num_samples" in prompt
+    assert "generated_samples" in prompt
+    assert "The output must be formatted:" in prompt
+    assert "Do not include any other text or break the schema in any way." in prompt
+    assert (
+        "Note how the output of this task is data to input into the system prompt"
+        in prompt
+    )
+
+
+def test_generate_sample_generation_prompt_structure_consistency():
+    """Test that the prompt structure is consistent between training and eval types"""
+    # Act
+    training_prompt = generate_sample_generation_prompt(gen_type="training")
+    eval_prompt = generate_sample_generation_prompt(gen_type="eval")
+
+    # Assert
+    # Both should have the same structure, just different goal descriptions
+    assert "## Task Description" in training_prompt
+    assert "## Task Description" in eval_prompt
+
+    # The main difference should be in the goal description
+    assert "train a large language model" in training_prompt
+    assert "evaluate a large language model" in eval_prompt
+    assert "generate training data" in training_prompt
+    assert "generate eval data" in eval_prompt
+
+    # Both should have the same core content
+    assert "Your job is to generate a list of potential inputs" in training_prompt
+    assert "Your job is to generate a list of potential inputs" in eval_prompt
+    assert "generated_samples" in training_prompt
+    assert "generated_samples" in eval_prompt
+
+
+def test_generate_sample_generation_prompt_with_none_guidance():
+    """Test generate_sample_generation_prompt with None guidance"""
+    # Act
+    prompt = generate_sample_generation_prompt(gen_type="training", guidance=None)
+
+    # Assert
+    assert isinstance(prompt, str)
+    assert (
+        "I want to train a large language model and you should help me generate training data for it."
+        in prompt
+    )
+    assert "## Specific Guidance" not in prompt
+    assert "The guidance is:" not in prompt

kiln_ai/adapters/eval/base_eval.py

@@ -2,8 +2,6 @@ import json
 from abc import abstractmethod
 from typing import Dict
 
-import jsonschema
-
 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
@@ -60,14 +58,13 @@ class BaseEval:
 
         run_adapter = adapter_for_task(
             self.target_task,
-            self.run_config.model_name,
-            ModelProviderName(self.run_config.model_provider_name),
+            self.run_config,
             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:
+        if self.target_task.input_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
@@ -122,7 +119,9 @@ class BaseEval:
                         property["minimum"] = 1
                         property["maximum"] = 5
                     else:
-                        property["enum"] = [1, 2, 3, 4, 5]
+                        property["type"] = "integer"
+                        property["minimum"] = 1
+                        property["maximum"] = 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."
@@ -137,6 +136,7 @@ class BaseEval:
                         )
                     else:
                         property["enum"] = ["pass", "fail"]
+                        property["type"] = "string"
                         property["description"] = (
                             f"{output_score.instruction}\n\nThe rating should be either 'pass' or 'fail'."
                         )
@@ -150,6 +150,7 @@ class BaseEval:
                         )
                     else:
                         property["enum"] = ["pass", "fail", "critical"]
+                        property["type"] = "string"
                         property["description"] = (
                             f"{output_score.instruction}\n\nThe rating should be either 'pass', 'fail', or 'critical' where critical a very severe failure."
                         )

kiln_ai/adapters/eval/eval_runner.py

@@ -8,7 +8,7 @@ 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
+from kiln_ai.datamodel.task_run import TaskRun, Usage
 from kiln_ai.utils.async_job_runner import AsyncJobRunner, Progress
 
 logger = logging.getLogger(__name__)
@@ -177,10 +177,12 @@ class EvalRunner:
             task_output: str | None = None
             scores: EvalScores | None = None
             intermediate_outputs: Dict[str, str] | None = None
+            task_run_usage: Usage | 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
+                task_run_usage = job.item.usage
             else:
                 # Task run eval, we invoke the task again to get a fresh output
                 (
@@ -189,6 +191,7 @@ class EvalRunner:
                     intermediate_outputs,
                 ) = await evaluator.run_task_and_eval(job.item.input)
                 task_output = result_task_run.output.output
+                task_run_usage = result_task_run.usage
 
             # Save the job result
             eval_run = EvalRun(
@@ -202,10 +205,14 @@ class EvalRunner:
                 input=job.item.input,
                 output=task_output,
                 intermediate_outputs=intermediate_outputs,
+                task_run_usage=task_run_usage,
             )
             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}")
+            logger.error(
+                f"Error running eval job for dataset item {job.item.id}: {e}",
+                exc_info=True,
+            )
             return False