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

kiln_ai/adapters/test_prompt_adaptors.py

@@ -6,13 +6,14 @@ from langchain_core.language_models.fake_chat_models import FakeListChatModel
 
 import kiln_ai.datamodel as datamodel
 from kiln_ai.adapters.adapter_registry import adapter_for_task
-from kiln_ai.adapters.langchain_adapters import LangchainAdapter
 from kiln_ai.adapters.ml_model_list import built_in_models
+from kiln_ai.adapters.model_adapters.langchain_adapters import LangchainAdapter
 from kiln_ai.adapters.ollama_tools import ollama_online
 from kiln_ai.adapters.prompt_builders import (
     BasePromptBuilder,
     SimpleChainOfThoughtPromptBuilder,
 )
+from kiln_ai.datamodel import PromptId
 
 
 def get_all_models_and_providers():
@@ -108,7 +109,11 @@ async def test_amazon_bedrock(tmp_path):
 async def test_mock(tmp_path):
     task = build_test_task(tmp_path)
     mockChatModel = FakeListChatModel(responses=["mock response"])
-    adapter = LangchainAdapter(task, custom_model=mockChatModel)
+    adapter = LangchainAdapter(
+        task,
+        custom_model=mockChatModel,
+        provider="ollama",
+    )
     run = await adapter.invoke("You are a mock, send me the response!")
     assert "mock response" in run.output.output
 
@@ -116,7 +121,7 @@ async def test_mock(tmp_path):
 async def test_mock_returning_run(tmp_path):
     task = build_test_task(tmp_path)
     mockChatModel = FakeListChatModel(responses=["mock response"])
-    adapter = LangchainAdapter(task, custom_model=mockChatModel)
+    adapter = LangchainAdapter(task, custom_model=mockChatModel, provider="ollama")
     run = await adapter.invoke("You are a mock, send me the response!")
     assert run.output.output == "mock response"
     assert run is not None
@@ -127,8 +132,8 @@ async def test_mock_returning_run(tmp_path):
     assert run.output.source.properties == {
         "adapter_name": "kiln_langchain_adapter",
         "model_name": "custom.langchain:unknown_model",
-        "model_provider": "custom.langchain:FakeListChatModel",
-        "prompt_builder_name": "simple_prompt_builder",
+        "model_provider": "ollama",
+        "prompt_id": "simple_prompt_builder",
     }
 
 
@@ -145,8 +150,9 @@ async def test_all_models_providers_plaintext(tmp_path, model_name, provider_nam
 @pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
 async def test_cot_prompt_builder(tmp_path, model_name, provider_name):
     task = build_test_task(tmp_path)
-    pb = SimpleChainOfThoughtPromptBuilder(task)
-    await run_simple_task(task, model_name, provider_name, pb)
+    await run_simple_task(
+        task, model_name, provider_name, "simple_chain_of_thought_prompt_builder"
+    )
 
 
 def build_test_task(tmp_path: Path):
@@ -182,20 +188,20 @@ async def run_simple_test(
     tmp_path: Path,
     model_name: str,
     provider: str | None = None,
-    prompt_builder: BasePromptBuilder | None = None,
+    prompt_id: PromptId | None = None,
 ):
     task = build_test_task(tmp_path)
-    return await run_simple_task(task, model_name, provider, prompt_builder)
+    return await run_simple_task(task, model_name, provider, prompt_id)
 
 
 async def run_simple_task(
     task: datamodel.Task,
     model_name: str,
     provider: str,
-    prompt_builder: BasePromptBuilder | None = None,
+    prompt_id: PromptId | None = None,
 ) -> datamodel.TaskRun:
     adapter = adapter_for_task(
-        task, model_name=model_name, provider=provider, prompt_builder=prompt_builder
+        task, model_name=model_name, provider=provider, prompt_id=prompt_id
     )
 
     run = await adapter.invoke(
@@ -208,13 +214,14 @@ async def run_simple_task(
     )
     assert "64" in run.output.output
     source_props = run.output.source.properties
-    assert source_props["adapter_name"] == "kiln_langchain_adapter"
+    assert source_props["adapter_name"] in [
+        "kiln_langchain_adapter",
+        "kiln_openai_compatible_adapter",
+    ]
     assert source_props["model_name"] == model_name
     assert source_props["model_provider"] == provider
-    expected_prompt_builder_name = (
-        prompt_builder.__class__.prompt_builder_name()
-        if prompt_builder
-        else "simple_prompt_builder"
-    )
-    assert source_props["prompt_builder_name"] == expected_prompt_builder_name
+    if prompt_id is None:
+        assert source_props["prompt_id"] == "simple_prompt_builder"
+    else:
+        assert source_props["prompt_id"] == prompt_id
     return run

kiln_ai/adapters/test_prompt_builders.py

@@ -1,37 +1,49 @@
 import json
+import logging
 
 import pytest
 
-from kiln_ai.adapters.base_adapter import AdapterInfo, BaseAdapter
+from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter
+from kiln_ai.adapters.model_adapters.test_structured_output import (
+    build_structured_output_test_task,
+)
 from kiln_ai.adapters.prompt_builders import (
     FewShotChainOfThoughtPromptBuilder,
     FewShotPromptBuilder,
+    FineTunePromptBuilder,
     MultiShotChainOfThoughtPromptBuilder,
     MultiShotPromptBuilder,
     RepairsPromptBuilder,
+    SavedPromptBuilder,
     SimpleChainOfThoughtPromptBuilder,
     SimplePromptBuilder,
+    TaskRunConfigPromptBuilder,
     chain_of_thought_prompt,
-    prompt_builder_from_ui_name,
+    prompt_builder_from_id,
 )
 from kiln_ai.adapters.test_prompt_adaptors import build_test_task
-from kiln_ai.adapters.test_structured_output import build_structured_output_test_task
 from kiln_ai.datamodel import (
     DataSource,
     DataSourceType,
+    Finetune,
+    FinetuneDataStrategy,
     Project,
+    Prompt,
     Task,
     TaskOutput,
     TaskOutputRating,
     TaskRun,
 )
+from kiln_ai.datamodel.task import RunConfigProperties, TaskRunConfig
+
+logger = logging.getLogger(__name__)
 
 
 def test_simple_prompt_builder(tmp_path):
     task = build_test_task(tmp_path)
     builder = SimplePromptBuilder(task=task)
     input = "two plus two"
-    prompt = builder.build_prompt()
+    prompt = builder.build_prompt(include_json_instructions=False)
     assert (
         "You are an assistant which performs math tasks provided in plain text."
         in prompt
@@ -50,19 +62,15 @@ class MockAdapter(BaseAdapter):
     def _run(self, input: str) -> str:
         return "mock response"
 
-    def adapter_info(self) -> AdapterInfo:
-        return AdapterInfo(
-            adapter_name="mock_adapter",
-            model_name="mock_model",
-            model_provider="mock_provider",
-        )
+    def adapter_name(self) -> str:
+        return "mock_adapter"
 
 
 def test_simple_prompt_builder_structured_output(tmp_path):
     task = build_structured_output_test_task(tmp_path)
     builder = SimplePromptBuilder(task=task)
     input = "Cows"
-    prompt = builder.build_prompt()
+    prompt = builder.build_prompt(include_json_instructions=False)
     assert "You are an assistant which tells a joke, given a subject." in prompt
 
     user_msg = builder.build_user_message(input)
@@ -70,6 +78,14 @@ def test_simple_prompt_builder_structured_output(tmp_path):
     assert input not in prompt
 
 
+def test_simple_prompt_builder_structured_input_non_ascii(tmp_path):
+    task = build_structured_output_test_task(tmp_path)
+    builder = SimplePromptBuilder(task=task)
+    input = {"key": "你好👋"}
+    user_msg = builder.build_user_message(input)
+    assert "你好👋" in user_msg
+
+
 @pytest.fixture
 def task_with_examples(tmp_path):
     # Create a project and task hierarchy
@@ -198,7 +214,7 @@ def task_with_examples(tmp_path):
 def test_multi_shot_prompt_builder(task_with_examples):
     # Verify the order of examples
     prompt_builder = MultiShotPromptBuilder(task=task_with_examples)
-    prompt = prompt_builder.build_prompt()
+    prompt = prompt_builder.build_prompt(include_json_instructions=False)
     assert "Why did the cow cross the road?" in prompt
     assert prompt.index("Why did the cow cross the road?") < prompt.index(
         "Why don't cats play poker in the jungle?"
@@ -239,14 +255,14 @@ def test_few_shot_prompt_builder(tmp_path):
     # Create 6 examples (2 repaired, 4 high-quality)
     for i in range(6):
         run = TaskRun(
-            input=f'{{"subject": "Subject {i+1}"}}',
+            input=f'{{"subject": "Subject {i + 1}"}}',
             input_source=DataSource(
                 type=DataSourceType.human,
                 properties={"created_by": "john_doe"},
             ),
             parent=task,
             output=TaskOutput(
-                output=f'{{"joke": "Joke Initial Output {i+1}"}}',
+                output=f'{{"joke": "Joke Initial Output {i + 1}"}}',
                 source=DataSource(
                     type=DataSourceType.human,
                     properties={"created_by": "john_doe"},
@@ -254,13 +270,12 @@ def test_few_shot_prompt_builder(tmp_path):
                 rating=TaskOutputRating(value=4 + (i % 2), reason="Good joke"),
             ),
         )
-        print("RATING", "Joke Initial Output ", i + 1, " - RATED:", 4 + (i % 2), "\n")
         if i < 2:
             run = run.model_copy(
                 update={
                     "repair_instructions": "Fix the joke",
                     "repaired_output": TaskOutput(
-                        output=f'{{"joke": "Repaired Joke {i+1}"}}',
+                        output=f'{{"joke": "Repaired Joke {i + 1}"}}',
                         source=DataSource(
                             type=DataSourceType.human,
                             properties={"created_by": "jane_doe"},
@@ -272,10 +287,10 @@ def test_few_shot_prompt_builder(tmp_path):
 
     # Check that only 4 examples are included
     prompt_builder = FewShotPromptBuilder(task=task)
-    prompt = prompt_builder.build_prompt()
+    prompt = prompt_builder.build_prompt(include_json_instructions=False)
     assert prompt.count("## Example") == 4
 
-    print("PROMPT", prompt)
+    logger.info("PROMPT: %s", prompt)
     # Verify the order of examples (2 repaired, then 2 highest-rated)
     assert "Repaired Joke 1" in prompt
     assert "Repaired Joke 2" in prompt
@@ -289,7 +304,7 @@ def test_few_shot_prompt_builder(tmp_path):
 
 def check_example_outputs(task: Task, count: int):
     prompt_builder = MultiShotPromptBuilder(task=task)
-    prompt = prompt_builder.build_prompt()
+    prompt = prompt_builder.build_prompt(include_json_instructions=False)
     assert "# Instruction" in prompt
     assert task.instruction in prompt
     if count == 0:
@@ -299,32 +314,89 @@ def check_example_outputs(task: Task, count: int):
         assert f"## Example {count}" in prompt
 
 
-def test_prompt_builder_name():
-    assert SimplePromptBuilder.prompt_builder_name() == "simple_prompt_builder"
-    assert MultiShotPromptBuilder.prompt_builder_name() == "multi_shot_prompt_builder"
-    assert RepairsPromptBuilder.prompt_builder_name() == "repairs_prompt_builder"
+def test_prompt_builder_from_id(task_with_examples):
+    task = task_with_examples
+    assert isinstance(
+        prompt_builder_from_id("simple_prompt_builder", task), SimplePromptBuilder
+    )
+    assert isinstance(
+        prompt_builder_from_id("few_shot_prompt_builder", task),
+        FewShotPromptBuilder,
+    )
+    assert isinstance(
+        prompt_builder_from_id("multi_shot_prompt_builder", task),
+        MultiShotPromptBuilder,
+    )
+    assert isinstance(
+        prompt_builder_from_id("repairs_prompt_builder", task),
+        RepairsPromptBuilder,
+    )
+    assert isinstance(
+        prompt_builder_from_id("simple_chain_of_thought_prompt_builder", task),
+        SimpleChainOfThoughtPromptBuilder,
+    )
+    assert isinstance(
+        prompt_builder_from_id("few_shot_chain_of_thought_prompt_builder", task),
+        FewShotChainOfThoughtPromptBuilder,
+    )
+    assert isinstance(
+        prompt_builder_from_id("multi_shot_chain_of_thought_prompt_builder", task),
+        MultiShotChainOfThoughtPromptBuilder,
+    )
+
+    with pytest.raises(ValueError, match="Unknown prompt generator: invalid_name"):
+        prompt_builder_from_id("invalid_name", task)
+
+    with pytest.raises(ValueError, match="Prompt ID not found: 123"):
+        prompt_builder_from_id("id::123", task)
 
+    with pytest.raises(
+        ValueError,
+        match="Invalid fine-tune ID format. Expected 'project_id::task_id::fine_tune_id'",
+    ):
+        prompt_builder_from_id("fine_tune_prompt::123", task)
 
-def test_prompt_builder_from_ui_name():
-    assert prompt_builder_from_ui_name("basic") == SimplePromptBuilder
-    assert prompt_builder_from_ui_name("few_shot") == FewShotPromptBuilder
-    assert prompt_builder_from_ui_name("many_shot") == MultiShotPromptBuilder
-    assert prompt_builder_from_ui_name("repairs") == RepairsPromptBuilder
-    assert (
-        prompt_builder_from_ui_name("simple_chain_of_thought")
-        == SimpleChainOfThoughtPromptBuilder
+    with pytest.raises(
+        ValueError,
+        match="Fine-tune ID not found",
+    ):
+        prompt_builder_from_id("fine_tune_prompt::123::456::789", task)
+
+    prompt = Prompt(
+        name="test_prompt_name",
+        prompt="test_prompt",
+        chain_of_thought_instructions="coti",
+        parent=task,
     )
-    assert (
-        prompt_builder_from_ui_name("few_shot_chain_of_thought")
-        == FewShotChainOfThoughtPromptBuilder
+    prompt.save_to_file()
+    pb = prompt_builder_from_id("id::" + prompt.id, task)
+    assert isinstance(pb, SavedPromptBuilder)
+    assert pb.prompt_id() == prompt.id
+    assert pb.build_prompt(include_json_instructions=False) == "test_prompt"
+    assert pb.chain_of_thought_prompt() == "coti"
+
+    finetune = Finetune(
+        name="test_finetune_name",
+        system_message="test_system_message",
+        thinking_instructions="test_thinking_instructions",
+        parent=task,
+        base_model_id="test_base_model_id",
+        dataset_split_id="asdf",
+        provider="test_provider",
+        data_strategy=FinetuneDataStrategy.final_and_intermediate,
     )
-    assert (
-        prompt_builder_from_ui_name("multi_shot_chain_of_thought")
-        == MultiShotChainOfThoughtPromptBuilder
+    finetune.save_to_file()
+    nested_fine_tune_id = (
+        task_with_examples.parent.id + "::" + task_with_examples.id + "::" + finetune.id
     )
-
-    with pytest.raises(ValueError, match="Unknown prompt builder: invalid_name"):
-        prompt_builder_from_ui_name("invalid_name")
+    pb = prompt_builder_from_id(
+        "fine_tune_prompt::" + nested_fine_tune_id,
+        task_with_examples,
+    )
+    assert isinstance(pb, FineTunePromptBuilder)
+    assert pb.prompt_id() == nested_fine_tune_id
+    assert pb.build_base_prompt() == "test_system_message"
+    assert pb.chain_of_thought_prompt() == "test_thinking_instructions"
 
 
 def test_example_count():
@@ -335,7 +407,7 @@ def test_example_count():
 def test_repair_multi_shot_prompt_builder(task_with_examples):
     # Verify the order of examples
     prompt_builder = RepairsPromptBuilder(task=task_with_examples)
-    prompt = prompt_builder.build_prompt()
+    prompt = prompt_builder.build_prompt(include_json_instructions=False)
     assert (
         'Repaired Output Which is Sufficient: {"joke": "Why did the cow cross the road? To get to the udder side!"}'
         in prompt
@@ -403,7 +475,7 @@ def test_build_prompt_for_ui(tmp_path):
     ui_prompt = simple_builder.build_prompt_for_ui()
 
     # Should match regular prompt since no chain of thought
-    assert ui_prompt == simple_builder.build_prompt()
+    assert ui_prompt == simple_builder.build_prompt(include_json_instructions=False)
     assert "# Thinking Instructions" not in ui_prompt
 
     # Test chain of thought prompt builder
@@ -411,7 +483,7 @@ def test_build_prompt_for_ui(tmp_path):
     ui_prompt_cot = cot_builder.build_prompt_for_ui()
 
     # Should include both base prompt and thinking instructions
-    assert cot_builder.build_prompt() in ui_prompt_cot
+    assert cot_builder.build_prompt(include_json_instructions=False) in ui_prompt_cot
     assert "# Thinking Instructions" in ui_prompt_cot
     assert "Think step by step" in ui_prompt_cot
 
@@ -423,6 +495,155 @@ def test_build_prompt_for_ui(tmp_path):
     custom_cot_builder = SimpleChainOfThoughtPromptBuilder(task=task_with_custom)
     ui_prompt_custom = custom_cot_builder.build_prompt_for_ui()
 
-    assert custom_cot_builder.build_prompt() in ui_prompt_custom
+    assert (
+        custom_cot_builder.build_prompt(include_json_instructions=False)
+        in ui_prompt_custom
+    )
     assert "# Thinking Instructions" in ui_prompt_custom
     assert custom_instruction in ui_prompt_custom
+
+
+def test_saved_prompt_builder(tmp_path):
+    task = build_test_task(tmp_path)
+
+    prompt = Prompt(
+        name="test_prompt_name",
+        prompt="test_prompt",
+        parent=task,
+    )
+    prompt.save_to_file()
+
+    builder = SavedPromptBuilder(task=task, prompt_id=prompt.id)
+    assert builder.build_prompt(include_json_instructions=False) == "test_prompt"
+    assert builder.chain_of_thought_prompt() is None
+    assert builder.build_prompt_for_ui() == "test_prompt"
+    assert builder.prompt_id() == prompt.id
+
+
+def test_saved_prompt_builder_with_chain_of_thought(tmp_path):
+    task = build_test_task(tmp_path)
+
+    prompt = Prompt(
+        name="test_prompt_name",
+        prompt="test_prompt",
+        chain_of_thought_instructions="Think step by step",
+        parent=task,
+    )
+    prompt.save_to_file()
+
+    builder = SavedPromptBuilder(task=task, prompt_id=prompt.id)
+    assert builder.build_prompt(include_json_instructions=False) == "test_prompt"
+    assert builder.chain_of_thought_prompt() == "Think step by step"
+    assert "Think step by step" in builder.build_prompt_for_ui()
+    assert builder.prompt_id() == prompt.id
+
+
+def test_saved_prompt_builder_not_found(tmp_path):
+    task = build_test_task(tmp_path)
+
+    with pytest.raises(ValueError, match="Prompt ID not found: 123"):
+        SavedPromptBuilder(task=task, prompt_id="123")
+
+
+def test_build_prompt_with_json_instructions(tmp_path):
+    task = build_test_task(tmp_path)
+    task = task.model_copy(
+        update={
+            "output_json_schema": json.dumps(
+                {
+                    "type": "object",
+                    "properties": {"result": {"type": "string"}},
+                    "required": ["result"],
+                }
+            )
+        }
+    )
+
+    builder = SimplePromptBuilder(task=task)
+
+    # Test without JSON instructions
+    prompt_without_json = builder.build_prompt(include_json_instructions=False)
+    assert "Format Instructions" not in prompt_without_json
+    assert (
+        "Return a JSON object conforming to the following schema:"
+        not in prompt_without_json
+    )
+    assert task.output_json_schema not in prompt_without_json
+
+    # Test with JSON instructions
+    prompt_with_json = builder.build_prompt(include_json_instructions=True)
+    assert "# Format Instructions" in prompt_with_json
+    assert (
+        "Return a JSON object conforming to the following schema:" in prompt_with_json
+    )
+    assert "```" in prompt_with_json
+    assert (
+        "{'type': 'object', 'properties': {'result': {'type': 'string'}}, 'required': ['result']}"
+        in prompt_with_json
+    )
+
+    # Verify base prompt is still included
+    assert task.instruction in prompt_with_json
+    for requirement in task.requirements:
+        assert requirement.instruction in prompt_with_json
+
+
+def test_task_run_config_prompt_builder(tmp_path):
+    task = build_test_task(tmp_path)
+
+    run_config = TaskRunConfig(
+        name="test_run_config",
+        parent=task,
+        run_config_properties=RunConfigProperties(
+            model_name="gpt-4",
+            model_provider_name="openai",
+            prompt_id="simple_prompt_builder",
+        ),
+        prompt=Prompt(
+            name="test prompt name",
+            prompt="test prompt content",
+            chain_of_thought_instructions="test step by step",
+        ),
+    )
+    run_config.save_to_file()
+
+    # Construct the eval prompt ID
+    run_config_prompt_id = (
+        f"task_run_config::{task.parent.id}::{task.id}::{run_config.id}"
+    )
+
+    # Test successful creation 2 ways: constructor and ID creation
+    builders = [
+        TaskRunConfigPromptBuilder(
+            task=task, run_config_prompt_id=run_config_prompt_id
+        ),
+        prompt_builder_from_id(run_config_prompt_id, task),
+    ]
+
+    for builder in builders:
+        assert (
+            builder.build_prompt(include_json_instructions=False)
+            == "test prompt content"
+        )
+        assert builder.chain_of_thought_prompt() == "test step by step"
+        assert builder.prompt_id() == run_config_prompt_id
+
+
+def test_task_run_config_prompt_builder_validation_errors(tmp_path):
+    task = build_test_task(tmp_path)
+
+    # Test invalid format
+    with pytest.raises(ValueError, match="Invalid task run config prompt ID"):
+        TaskRunConfigPromptBuilder(
+            task=task, run_config_prompt_id="task_run_config::wrong::format"
+        )
+
+    # Test task ID mismatch
+    wrong_task_id = f"task_run_config::{task.parent.id}::wrong_task_id::config_id"
+    with pytest.raises(ValueError, match="Task ID mismatch"):
+        TaskRunConfigPromptBuilder(task=task, run_config_prompt_id=wrong_task_id)
+
+    # Test eval not found
+    nonexistent_eval = f"task_run_config::{task.parent.id}::{task.id}::nonexistent_id"
+    with pytest.raises(ValueError, match="Task run config ID not found"):
+        TaskRunConfigPromptBuilder(task=task, run_config_prompt_id=nonexistent_eval)