kiln-ai 0.5.5__tar.gz → 0.6.0__tar.gz

{kiln_ai-0.5.5 → kiln_ai-0.6.0}/PKG-INFO

@@ -1,13 +1,12 @@
 Metadata-Version: 2.3
 Name: kiln-ai
-Version: 0.5.5
+Version: 0.6.0
 Summary: Kiln AI
 Project-URL: Homepage, https://getkiln.ai
 Project-URL: Repository, https://github.com/Kiln-AI/kiln
 Project-URL: Documentation, https://kiln-ai.github.io/Kiln/kiln_core_docs/kiln_ai.html
 Project-URL: Issues, https://github.com/Kiln-AI/kiln/issues
 Author-email: "Steve Cosman, Chesterfield Laboratories Inc" <scosman@users.noreply.github.com>
-License-File: LICENSE.txt
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3.10

{kiln_ai-0.5.5 → kiln_ai-0.6.0}/kiln_ai/adapters/base_adapter.py

@@ -24,6 +24,12 @@ class AdapterInfo:
     prompt_builder_name: str
 
 
+@dataclass
+class RunOutput:
+    output: Dict | str
+    intermediate_outputs: Dict[str, str] | None
+
+
 class BaseAdapter(metaclass=ABCMeta):
     """Base class for AI model adapters that handle task execution.
 
@@ -36,22 +42,6 @@ class BaseAdapter(metaclass=ABCMeta):
         kiln_task (Task): The task configuration and metadata
         output_schema (dict | None): JSON schema for validating structured outputs
         input_schema (dict | None): JSON schema for validating structured inputs
-
-    Example:
-        ```python
-        class CustomAdapter(BaseAdapter):
-            async def _run(self, input: Dict | str) -> Dict | str:
-                # Implementation for specific model
-                pass
-
-            def adapter_info(self) -> AdapterInfo:
-                return AdapterInfo(
-                    adapter_name="custom",
-                    model_name="model-1",
-                    model_provider="provider",
-                    prompt_builder_name="simple"
-                )
-        ```
     """
 
     def __init__(
@@ -85,21 +75,23 @@ class BaseAdapter(metaclass=ABCMeta):
             validate_schema(input, self.input_schema)
 
         # Run
-        result = await self._run(input)
+        run_output = await self._run(input)
 
         # validate output
         if self.output_schema is not None:
-            if not isinstance(result, dict):
-                raise RuntimeError(f"structured response is not a dict: {result}")
-            validate_schema(result, self.output_schema)
+            if not isinstance(run_output.output, dict):
+                raise RuntimeError(
+                    f"structured response is not a dict: {run_output.output}"
+                )
+            validate_schema(run_output.output, self.output_schema)
         else:
-            if not isinstance(result, str):
+            if not isinstance(run_output.output, str):
                 raise RuntimeError(
-                    f"response is not a string for non-structured task: {result}"
+                    f"response is not a string for non-structured task: {run_output.output}"
                 )
 
         # Generate the run and output
-        run = self.generate_run(input, input_source, result)
+        run = self.generate_run(input, input_source, run_output)
 
         # Save the run if configured to do so, and we have a path to save to
         if Config.shared().autosave_runs and self.kiln_task.path is not None:
@@ -118,27 +110,23 @@ class BaseAdapter(metaclass=ABCMeta):
         pass
 
     @abstractmethod
-    async def _run(self, input: Dict | str) -> Dict | str:
+    async def _run(self, input: Dict | str) -> RunOutput:
         pass
 
     def build_prompt(self) -> str:
-        prompt = self.prompt_builder.build_prompt()
-        adapter_instructions = self.adapter_specific_instructions()
-        if adapter_instructions is not None:
-            prompt += f"# Format Instructions\n\n{adapter_instructions}\n\n"
-        return prompt
-
-    # override for adapter specific instructions (e.g. tool calling, json format, etc)
-    def adapter_specific_instructions(self) -> str | None:
-        return None
+        return self.prompt_builder.build_prompt()
 
     # create a run and task output
     def generate_run(
-        self, input: Dict | str, input_source: DataSource | None, output: Dict | str
+        self, input: Dict | str, input_source: DataSource | None, run_output: RunOutput
     ) -> TaskRun:
         # Convert input and output to JSON strings if they are dictionaries
         input_str = json.dumps(input) if isinstance(input, dict) else input
-        output_str = json.dumps(output) if isinstance(output, dict) else output
+        output_str = (
+            json.dumps(run_output.output)
+            if isinstance(run_output.output, dict)
+            else run_output.output
+        )
 
         # If no input source is provided, use the human data source
         if input_source is None:
@@ -159,6 +147,7 @@ class BaseAdapter(metaclass=ABCMeta):
                     properties=self._properties_for_task_output(),
                 ),
             ),
+            intermediate_outputs=run_output.intermediate_outputs,
         )
 
         exclude_fields = {

kiln_ai-0.6.0/kiln_ai/adapters/data_gen/data_gen_prompts.py

@@ -0,0 +1,73 @@
+# The contents of this file are adapted from the promptwrite library (https://github.com/StacklokLabs/promptwright),
+# which was adapted from the pluto library (https://github.com/redotvideo/pluto).
+# These libraries are licensed under the Apache License 2.0. Any modifications
+# are licensed under the kiln AI Core license (MIT at time of writing). See /libs/core/LICENSE.txt for details.
+
+
+TREE_GENERATION_PROMPT = """I want to train a large language model and I am using another, bigger large language model to generate training data for this. However, if we always ask the bigger model to generate training data with the same prompt, it will end up generating very repetitive training samples. Therefore, we will slightly modify our prompt for each sampling procedure according to some aspects. For instance, when asking the model to generate news articles, we could modify the prompt to let the model tell news articles about particular topics, such as business or politics. To further generate training data, we will do this recursively, and generate submodifications to the prompt. For instance, within the domain of business, we could adapt the prompt to generate news about the stock market or business scandals, and within politics, we could ask the model to generate articles for subtopics like elections or climate policy. We do this recursively, and therefore, we get a tree-like structure of topics.
+Your job is the following: I will give you a path of nodes down the topic tree - you should then come up with a list of new subtopics for this given node and return it as a python list. Here are a few examples of what your outputs should look like, related to the news example I just gave you:
+
+Example 1:
+node path: "News Topics" -> "Sports" -> "Football"
+desired number of subtopics: 5
+subtopics: ["College Football", "Football Stadiums", "Health Consequences Football", "Seattle Seahawks", "Football Sponsorships"]
+
+Example 2:
+node path: "News Topics" -> "Entertainment" -> "Movies" -> "Star Portraits"
+desired number of subtopics: 8
+subtopics: ["Tom Hanks", "Meryl Streep", "Leonardo DiCaprio", "Jennifer Lawrence", "Denzel Washington", "Charlize Theron", "Robert Downey Jr.", "Emma Stone"]
+
+
+Here are three new examples, this time for generating smalltalk topics for a friendly chat assistant:
+
+Example 1:
+node path: "Small Talk Topics"
+desired number of subtopics: 7
+subtopics: ["Weather", "Weekend Plans", "Hobbies", "Family", "Books", "Food", "Music"]
+
+Example 2:
+node path: "Small Talk Topics" -> "Family"
+desired number of subtopics: 5
+subtopics: ["Parents", "Grandparents", "Siblings", "Family Traditions", "Family Vacations"]
+
+Example 3:
+node path: "Small Talk Topics" -> "Hobbies" -> "Cooking"
+desired number of subtopics: 6
+subtopics: ["Recipes", "Asian Food", "Favourite Dishes", "Cookbooks", "Kitchen Gadgets", "Vegan Cooking"]
+
+The user message will contain the following:
+ - The system prompt for the model we want to train as system_prompt.
+ - The node path as node_path. It will be formated as a list of strings from most general to most specific. For example, the node_path for Example 3 above would be ["Small Talk Topics", "Hobbies", "Cooking"]. If empty, the node path is the root node.
+ - The desired number of subtopics for this node as num_subtopics. Return exactly this number of subtopics.
+ - Optionally, it may contain human_guidance, which is a string that contains additional instructions for how to generate the subtopics.
+ - Optionally, it may contain existing_topics, which is a list of subtopics that already exist at this node. You should not generate subtopics that are in this list.
+
+
+When generating subtopics, remain somewhat vague. Things can only be tangentially related and they don't have to be interpreted in a single way. Importantly, make sure that the subtopics fit the system prompt.
+"""
+
+
+SAMPLE_GENERATION_PROMPT = """I want to train a large language model and you should help me generate training data for it. 
+
+Your job is to generate a list of potential inputs to the provided system prompt. They should be diverse and relevant to the system prompt, and the topic if provided.
+
+In the user message we'll provide the following:
+ - The system prompt as system_prompt
+ - A potential topic to generate samples for. This will be a list of strings from most general to most specific. For example, the topic ["Small Talk Topics", "Hobbies", "Cooking"] would represent the topic "Cooking" in the "Hobbies" category of "Small Talk Topics". The list may be empty, in which case you should generate samples using the system prompt alone.
+ - The number of samples to generate as num_samples. If greater than 1, generate a range of samples that are diverse and relevant to the system prompt, and the topic if provided.
+ - The user message may optionally contain human_guidance, which is a string that contains additional instructions for how to generate the samples.
+
+The output must be formatted:
+ - in the provided structured format, as an object with a single property "generated_samples" that maps to a list of generated samples that would be inputs to the provided system prompt.
+ - With the correct number of samples (num_samples).
+ - Do not include any other text or break the schema in any way.
+
+Example inputs:
+ - system_prompt: "You are an assistant that classifies the tone of a tweet. You should output one of the following labels: 'positive', 'negative', 'neutral'."
+ - topic: ["Technology", "New iPhone Event"]
+ - num_samples: 2
+Example output: {"generated_samples": ["New iPhone looks amazing! I need that camera.", "Another boring event from Apple.", "New iPhone looks interesting, but I'm waiting for reviews."]}
+
+
+Note how the output of this task is data to input to the system prompt, not the expected output of the system prompt.
+"""

kiln_ai-0.6.0/kiln_ai/adapters/data_gen/data_gen_task.py

@@ -0,0 +1,117 @@
+import json
+
+from kiln_ai.adapters.prompt_builders import SimplePromptBuilder
+from kiln_ai.datamodel import Project, Task
+from pydantic import BaseModel
+
+from .data_gen_prompts import (
+    SAMPLE_GENERATION_PROMPT,
+    TREE_GENERATION_PROMPT,
+)
+
+
+class DataGenCategoriesTaskInput(BaseModel):
+    node_path: list[str]
+    system_prompt: str
+    num_subtopics: int
+    human_guidance: str | None = None
+    existing_topics: list[str] | None = None
+
+    @classmethod
+    def from_task(
+        cls,
+        task: Task,
+        node_path: list[str] = [],
+        num_subtopics: int = 6,
+        human_guidance: str | None = None,
+        existing_topics: list[str] | None = None,
+    ) -> "DataGenCategoriesTaskInput":
+        prompt_builder = SimplePromptBuilder(task=task)
+        return cls(
+            node_path=node_path,
+            num_subtopics=num_subtopics,
+            human_guidance=human_guidance,
+            existing_topics=existing_topics,
+            system_prompt=prompt_builder.build_prompt(),
+        )
+
+
+class DataGenCategoriesTaskOutput(BaseModel):
+    subtopics: list[str]
+
+
+class DataGenCategoriesTask(Task, parent_of={}):
+    def __init__(self):
+        # Keep the typechecker happy. TODO: shouldn't need this or parent_of above.
+        tmp_project = Project(name="DataGen")
+        super().__init__(
+            name="DataGen",
+            parent=tmp_project,
+            description="A task which generates synthetic data categories, which in turn are used to generate training data for a model to learn from.",
+            instruction=TREE_GENERATION_PROMPT,
+            input_json_schema=json.dumps(
+                DataGenCategoriesTaskInput.model_json_schema()
+            ),
+            output_json_schema=json.dumps(
+                DataGenCategoriesTaskOutput.model_json_schema()
+            ),
+        )
+
+
+class DataGenSampleTaskInput(BaseModel):
+    topic: list[str]
+    system_prompt: str
+    num_samples: int
+    human_guidance: str | None = None
+
+    @classmethod
+    def from_task(
+        cls,
+        task: Task,
+        topic: list[str] = [],
+        num_samples: int = 8,
+        human_guidance: str | None = None,
+    ) -> "DataGenSampleTaskInput":
+        prompt_builder = SimplePromptBuilder(task=task)
+        return cls(
+            topic=topic,
+            num_samples=num_samples,
+            human_guidance=human_guidance,
+            system_prompt=prompt_builder.build_prompt(),
+        )
+
+
+def list_json_schema_for_task(task: Task) -> str:
+    if task.input_json_schema:
+        items_schema = json.loads(task.input_json_schema)
+    else:
+        items_schema = {"type": "string"}
+
+    list_schema = {
+        "type": "array",
+        "items": items_schema,
+    }
+
+    top_level_schema = {
+        "type": "object",
+        "properties": {
+            "generated_samples": list_schema,
+        },
+        "required": ["generated_samples"],
+    }
+
+    return json.dumps(top_level_schema)
+
+
+class DataGenSampleTask(Task, parent_of={}):
+    def __init__(self, target_task: Task, num_samples: int = 8):
+        # Keep the typechecker happy. TODO: shouldn't need this or parent_of above.
+        tmp_project = Project(name="DataGenSample")
+        super().__init__(
+            name="DataGenSample",
+            parent=tmp_project,
+            description="A task which generates synthetic data samples for a given topic (and optional subtopic).",
+            instruction=SAMPLE_GENERATION_PROMPT,
+            input_json_schema=json.dumps(DataGenSampleTaskInput.model_json_schema()),
+            output_json_schema=list_json_schema_for_task(target_task),
+        )

kiln_ai-0.6.0/kiln_ai/adapters/data_gen/test_data_gen_task.py

@@ -0,0 +1,292 @@
+import json
+
+import pytest
+from kiln_ai.adapters.data_gen.data_gen_task import (
+    DataGenCategoriesTask,
+    DataGenCategoriesTaskInput,
+    DataGenCategoriesTaskOutput,
+    DataGenSampleTask,
+    DataGenSampleTaskInput,
+    list_json_schema_for_task,
+)
+from kiln_ai.adapters.langchain_adapters import LangChainPromptAdapter
+from kiln_ai.adapters.ml_model_list 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
+
+
+@pytest.fixture
+def base_task():
+    project = Project(name="TestProject")
+    return Task(
+        name="Cowboy Speaker",
+        parent=project,
+        description="Reply like a cowboy",
+        instruction="Reply like a cowboy",
+        requirements=[],
+    )
+
+
+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
+
+
+def test_data_gen_categories_task_input_default_values(base_task):
+    # Act
+    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 == []
+
+
+def test_data_gen_categories_task_initialization():
+    # Act
+    task = DataGenCategoriesTask()
+
+    # Assert
+    assert task.name == "DataGen"
+    assert isinstance(task.parent, Project)
+    assert task.description is not None
+    assert task.instruction is not None
+    assert isinstance(task.input_json_schema, str)
+    assert isinstance(task.output_json_schema, str)
+
+
+def test_data_gen_categories_task_schemas():
+    # Act
+    task = DataGenCategoriesTask()
+
+    # Assert
+    input_schema = json.loads(task.input_json_schema)
+    output_schema = json.loads(task.output_json_schema)
+
+    assert isinstance(input_schema, dict)
+    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 set(input_schema["required"]) == {
+        "node_path",
+        "num_subtopics",
+        "system_prompt",
+    }
+
+
+@pytest.mark.paid
+@pytest.mark.ollama
+@pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
+async def test_data_gen_all_models_providers(
+    tmp_path, model_name, provider_name, base_task
+):
+    _, provider = get_model_and_provider(model_name, provider_name)
+    if not provider.supports_data_gen:
+        # pass if the model doesn't support data gen (testing the support flag is part of this)
+        return
+
+    data_gen_task = DataGenCategoriesTask()
+    data_gen_input = DataGenCategoriesTaskInput.from_task(base_task, num_subtopics=6)
+
+    adapter = LangChainPromptAdapter(
+        data_gen_task,
+        model_name=model_name,
+        provider=provider_name,
+    )
+
+    input_dict = data_gen_input.model_dump()
+    run = await adapter.invoke(input_dict)
+    parsed_output = DataGenCategoriesTaskOutput.model_validate_json(run.output.output)
+    assert len(parsed_output.subtopics) == 6
+    for subtopic in parsed_output.subtopics:
+        assert isinstance(subtopic, str)
+
+
+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
+
+
+def test_data_gen_sample_task_input_default_values(base_task):
+    # Act
+    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 == []
+
+
+def test_data_gen_sample_task_initialization(base_task):
+    # Act
+    task = DataGenSampleTask(target_task=base_task)
+
+    # Assert
+    assert task.name == "DataGenSample"
+    assert isinstance(task.parent, Project)
+    assert task.description is not None
+    assert task.instruction is not None
+
+    input_schema = json.loads(task.input_json_schema)
+    output_schema = json.loads(task.output_json_schema)
+
+    assert isinstance(input_schema, dict)
+    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 set(input_schema["required"]) == {
+        "topic",
+        "num_samples",
+        "system_prompt",
+    }
+
+
+def test_list_json_schema_for_task_with_output_schema(base_task):
+    # Arrange
+    base_task.input_json_schema = json.dumps(
+        {
+            "type": "object",
+            "properties": {"name": {"type": "string"}, "age": {"type": "integer"}},
+        }
+    )
+
+    # Act
+    schema = list_json_schema_for_task(base_task)
+    parsed_schema = json.loads(schema)
+
+    # Assert
+    assert parsed_schema["type"] == "object"
+    generated_samples_schema = parsed_schema["properties"]["generated_samples"]
+    assert generated_samples_schema["type"] == "array"
+    assert generated_samples_schema["items"]["type"] == "object"
+    assert generated_samples_schema["items"]["properties"]["name"]["type"] == "string"
+    assert generated_samples_schema["items"]["properties"]["age"]["type"] == "integer"
+
+
+def test_list_json_schema_for_task_without_output_schema(base_task):
+    # Arrange
+    base_task.output_json_schema = None
+
+    # Act
+    schema = list_json_schema_for_task(base_task)
+    parsed_schema = json.loads(schema)
+
+    # Assert
+    assert parsed_schema["type"] == "object"
+    assert parsed_schema["properties"]["generated_samples"]["type"] == "array"
+    assert parsed_schema["properties"]["generated_samples"]["items"]["type"] == "string"
+
+
+@pytest.mark.paid
+@pytest.mark.ollama
+@pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
+async def test_data_gen_sample_all_models_providers(
+    tmp_path, model_name, provider_name, base_task
+):
+    _, provider = get_model_and_provider(model_name, provider_name)
+    if not provider.supports_data_gen:
+        # pass if the model doesn't support data gen (testing the support flag is part of this)
+        return
+
+    data_gen_task = DataGenSampleTask(target_task=base_task)
+    data_gen_input = DataGenSampleTaskInput.from_task(
+        base_task, topic=["riding horses"], num_samples=4
+    )
+
+    adapter = LangChainPromptAdapter(
+        data_gen_task,
+        model_name=model_name,
+        provider=provider_name,
+    )
+
+    input_dict = data_gen_input.model_dump()
+    run = await adapter.invoke(input_dict)
+    parsed_output = json.loads(run.output.output)
+    samples = parsed_output["generated_samples"]
+    assert len(samples) == 4
+    for sample in samples:
+        assert isinstance(sample, str)
+
+
+@pytest.mark.paid
+@pytest.mark.ollama
+@pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
+async def test_data_gen_sample_all_models_providers_with_structured_output(
+    tmp_path, model_name, provider_name, base_task
+):
+    base_task.output_json_schema = json.dumps(
+        {
+            "type": "object",
+            "properties": {
+                "opening": {"type": "string"},
+                "closing": {"type": "string"},
+            },
+            "required": ["opening", "closing"],
+        }
+    )
+
+    _, provider = get_model_and_provider(model_name, provider_name)
+    if not provider.supports_data_gen:
+        # pass if the model doesn't support data gen (testing the support flag is part of this)
+        return
+
+    data_gen_task = DataGenSampleTask(target_task=base_task)
+    data_gen_input = DataGenSampleTaskInput.from_task(
+        base_task, topic=["riding horses"], num_samples=4
+    )
+
+    adapter = LangChainPromptAdapter(
+        data_gen_task,
+        model_name=model_name,
+        provider=provider_name,
+    )
+
+    input_dict = data_gen_input.model_dump()
+    run = await adapter.invoke(input_dict)
+    parsed_output = json.loads(run.output.output)
+    samples = parsed_output["generated_samples"]
+    assert len(samples) == 4
+    for sample in samples:
+        assert isinstance(sample, dict)
+        assert "opening" in sample
+        assert "closing" in sample
+        assert isinstance(sample["opening"], str)
+        assert isinstance(sample["closing"], str)