kiln-ai 0.17.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,
@@ -33,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):
@@ -56,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"
@@ -72,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)
@@ -86,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",
     }
 
 
@@ -106,7 +113,7 @@ 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(
@@ -131,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):
@@ -154,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)
@@ -176,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",
     }
 
 
@@ -332,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

@@ -7,12 +7,7 @@ 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_with_value_error
-from kiln_ai.datamodel.task import (
-    RunConfig,
-    RunConfigProperties,
-    TaskOutputRatingType,
-    TaskRun,
-)
+from kiln_ai.datamodel.task import RunConfig, TaskOutputRatingType, TaskRun
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 
@@ -124,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."
@@ -139,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'."
                         )
@@ -152,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,6 +205,7 @@ class EvalRunner:
                 input=job.item.input,
                 output=task_output,
                 intermediate_outputs=intermediate_outputs,
+                task_run_usage=task_run_usage,
             )
             eval_run.save_to_file()
 

kiln_ai/adapters/eval/g_eval.py

@@ -102,6 +102,18 @@ class GEval(BaseEval):
 
         self.geval_task = GEvalTask(eval_config)
 
+    def generate_run_description(self, eval_input: str, eval_output: str) -> str:
+        return f"""The model was given the following input for the task: 
+<eval_data>
+{eval_input}
+</eval_data>
+
+The model produced the following output for the task:
+<eval_data>
+{eval_output}
+</eval_data>
+"""
+
     async def run_eval(
         self, task_run: TaskRun
     ) -> tuple[EvalScores, Dict[str, str] | None]:
@@ -145,19 +157,12 @@ class GEval(BaseEval):
             ),
         )
 
-        input = f"""The model was given the following input for the task: 
-<eval_data>
-{task_run.input}
-</eval_data>
-
-The model produced the following output for the task:
-<eval_data>
-{task_run.output}
-</eval_data>
-"""
+        run_description = self.generate_run_description(
+            task_run.input, task_run.output.output
+        )
 
         # We don't need the run, but invoke_returning_run_output() runs validations for us over _run()
-        _, run_output = await adapter.invoke_returning_run_output(input)
+        _, run_output = await adapter.invoke_returning_run_output(run_description)
 
         if self.eval_config.config_type == EvalConfigType.llm_as_judge:
             return self.build_llm_as_judge_score(
@@ -310,7 +315,7 @@ The model produced the following output for the task:
         """
         primary_token_score = self.score_from_token_string(token_logprob.token)
         # check this is a real rating token, it could just be the ": ", "," or whitespace
-        if not primary_token_score:
+        if primary_token_score is None:
             return None
 
         total_score = 0.0

kiln_ai/adapters/eval/test_base_eval.py

@@ -43,7 +43,9 @@ def test_score_schema_five_star():
 
     # Check score property, and that it's an enum of 1-5
     score_prop = schema["properties"]["quality_score"]
-    assert score_prop["enum"] == [1, 2, 3, 4, 5]
+    assert score_prop["type"] == "integer"
+    assert score_prop["minimum"] == 1
+    assert score_prop["maximum"] == 5
     assert "Quality Score" in score_prop["title"]
     assert "Rate the quality" in score_prop["description"]
     assert "between 1 and 5" in score_prop["description"]
@@ -51,7 +53,9 @@ def test_score_schema_five_star():
     # Check overall rating property, and that it's an enum of 1-5
     assert "overall_rating" in schema["properties"]
     overall = schema["properties"]["overall_rating"]
-    assert overall["enum"] == [1, 2, 3, 4, 5]
+    assert overall["type"] == "integer"
+    assert overall["minimum"] == 1
+    assert overall["maximum"] == 5
     assert "Overall Rating" in overall["title"]
     assert "The overall rating for the task output" in overall["description"]
     assert "between 1 and 5" in overall["description"]
@@ -127,6 +131,7 @@ def test_score_schema_pass_fail():
     schema = json.loads(schema_str)
 
     score_prop = schema["properties"]["pass_fail_test"]
+    assert score_prop["type"] == "string"
     assert score_prop["enum"] == ["pass", "fail"]
     assert "Pass Fail Test" in score_prop["title"]
     assert "Check if it passes" in score_prop["description"]
@@ -173,6 +178,7 @@ def test_score_schema_pass_fail_critical():
     score_prop = schema["properties"]["critical_test"]
     assert "enum" in score_prop
     assert score_prop["enum"] == ["pass", "fail", "critical"]
+    assert score_prop["type"] == "string"
     assert "'pass', 'fail', or 'critical'" in score_prop["description"]
 
     assert schema["properties"]["overall_rating"] is not None