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

kiln_ai/adapters/test_adapter_registry.py

@@ -0,0 +1,177 @@
+from unittest.mock import patch
+
+import pytest
+
+from kiln_ai import datamodel
+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.langchain_adapters import LangchainAdapter
+from kiln_ai.adapters.model_adapters.openai_model_adapter import OpenAICompatibleAdapter
+from kiln_ai.adapters.prompt_builders import BasePromptBuilder
+from kiln_ai.adapters.provider_tools import kiln_model_provider_from
+
+
+@pytest.fixture
+def mock_config():
+    with patch("kiln_ai.adapters.adapter_registry.Config") as mock:
+        mock.shared.return_value.open_ai_api_key = "test-openai-key"
+        mock.shared.return_value.open_router_api_key = "test-openrouter-key"
+        yield mock
+
+
+@pytest.fixture
+def basic_task():
+    return datamodel.Task(
+        task_id="test-task",
+        task_type="test",
+        input_text="test input",
+        name="test-task",
+        instruction="test-task",
+    )
+
+
+@pytest.fixture
+def mock_finetune_from_id():
+    with patch("kiln_ai.adapters.provider_tools.finetune_from_id") as mock:
+        mock.return_value.provider = ModelProviderName.openai
+        mock.return_value.fine_tune_model_id = "test-model"
+        yield mock
+
+
+def test_openai_adapter_creation(mock_config, basic_task):
+    adapter = adapter_for_task(
+        kiln_task=basic_task, model_name="gpt-4", provider=ModelProviderName.openai
+    )
+
+    assert isinstance(adapter, OpenAICompatibleAdapter)
+    assert adapter.config.model_name == "gpt-4"
+    assert adapter.config.api_key == "test-openai-key"
+    assert adapter.config.provider_name == ModelProviderName.openai
+    assert adapter.config.base_url is None  # OpenAI url is default
+    assert adapter.config.default_headers is None
+
+
+def test_openrouter_adapter_creation(mock_config, basic_task):
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        model_name="anthropic/claude-3-opus",
+        provider=ModelProviderName.openrouter,
+    )
+
+    assert isinstance(adapter, OpenAICompatibleAdapter)
+    assert adapter.config.model_name == "anthropic/claude-3-opus"
+    assert adapter.config.api_key == "test-openrouter-key"
+    assert adapter.config.provider_name == ModelProviderName.openrouter
+    assert adapter.config.base_url == "https://openrouter.ai/api/v1"
+    assert adapter.config.default_headers == {
+        "HTTP-Referer": "https://getkiln.ai/openrouter",
+        "X-Title": "KilnAI",
+    }
+
+
+@pytest.mark.parametrize(
+    "provider",
+    [
+        ModelProviderName.groq,
+        ModelProviderName.amazon_bedrock,
+        ModelProviderName.ollama,
+        ModelProviderName.fireworks_ai,
+    ],
+)
+def test_langchain_adapter_creation(mock_config, basic_task, provider):
+    adapter = adapter_for_task(
+        kiln_task=basic_task, model_name="test-model", provider=provider
+    )
+
+    assert isinstance(adapter, LangchainAdapter)
+    assert adapter.model_name == "test-model"
+
+
+# TODO should run for all cases
+def test_custom_prompt_builder(mock_config, basic_task):
+    class TestPromptBuilder(BasePromptBuilder):
+        def build_base_prompt(self, kiln_task) -> str:
+            return "test-prompt"
+
+    prompt_builder = TestPromptBuilder(basic_task)
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        model_name="gpt-4",
+        provider=ModelProviderName.openai,
+        prompt_builder=prompt_builder,
+    )
+
+    assert adapter.prompt_builder == prompt_builder
+
+
+# TODO should run for all cases
+def test_tags_passed_through(mock_config, basic_task):
+    tags = ["test-tag-1", "test-tag-2"]
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        model_name="gpt-4",
+        provider=ModelProviderName.openai,
+        tags=tags,
+    )
+
+    assert adapter.default_tags == tags
+
+
+def test_invalid_provider(mock_config, basic_task):
+    with pytest.raises(ValueError, match="Unhandled enum value"):
+        adapter_for_task(
+            kiln_task=basic_task, model_name="test-model", provider="invalid"
+        )
+
+
+@patch("kiln_ai.adapters.adapter_registry.openai_compatible_config")
+def test_openai_compatible_adapter(mock_compatible_config, mock_config, basic_task):
+    mock_compatible_config.return_value.model_name = "test-model"
+    mock_compatible_config.return_value.api_key = "test-key"
+    mock_compatible_config.return_value.base_url = "https://test.com/v1"
+
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        model_name="provider::test-model",
+        provider=ModelProviderName.openai_compatible,
+    )
+
+    assert isinstance(adapter, OpenAICompatibleAdapter)
+    mock_compatible_config.assert_called_once_with("provider::test-model")
+    assert adapter.config.model_name == "test-model"
+    assert adapter.config.api_key == "test-key"
+    assert adapter.config.base_url == "https://test.com/v1"
+
+
+def test_custom_openai_compatible_provider(mock_config, basic_task):
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        model_name="openai::test-model",
+        provider=ModelProviderName.kiln_custom_registry,
+    )
+
+    assert isinstance(adapter, OpenAICompatibleAdapter)
+    assert adapter.config.model_name == "openai::test-model"
+    assert adapter.config.api_key == "test-openai-key"
+    assert adapter.config.base_url is None  # openai is none
+    assert adapter.config.provider_name == ModelProviderName.kiln_custom_registry
+
+
+async def test_fine_tune_provider(mock_config, basic_task, mock_finetune_from_id):
+    adapter = adapter_for_task(
+        kiln_task=basic_task,
+        model_name="proj::task::tune",
+        provider=ModelProviderName.kiln_fine_tune,
+    )
+
+    mock_finetune_from_id.assert_called_once_with("proj::task::tune")
+    assert isinstance(adapter, OpenAICompatibleAdapter)
+    assert adapter.config.provider_name == ModelProviderName.kiln_fine_tune
+    # Kiln model name here, but the underlying openai model id below
+    assert adapter.config.model_name == "proj::task::tune"
+
+    provider = kiln_model_provider_from(
+        "proj::task::tune", provider_name=ModelProviderName.kiln_fine_tune
+    )
+    # The actual model name from the fine tune object
+    assert provider.provider_options["model"] == "test-model"

kiln_ai/adapters/test_generate_docs.py

@@ -0,0 +1,69 @@
+from typing import List
+
+import pytest
+
+from libs.core.kiln_ai.adapters.ml_model_list import (
+    KilnModelProvider,
+    built_in_models,
+)
+from libs.core.kiln_ai.adapters.provider_tools import provider_name_from_id
+
+
+def _all_providers_support(providers: List[KilnModelProvider], attribute: str) -> bool:
+    """Check if all providers support a given feature"""
+    return all(getattr(provider, attribute) for provider in providers)
+
+
+def _any_providers_support(providers: List[KilnModelProvider], attribute: str) -> bool:
+    """Check if any providers support a given feature"""
+    return any(getattr(provider, attribute) for provider in providers)
+
+
+def _get_support_status(providers: List[KilnModelProvider], attribute: str) -> str:
+    """Get the support status for a feature"""
+    if _all_providers_support(providers, attribute):
+        return "✅︎"
+    elif _any_providers_support(providers, attribute):
+        return "✅︎ (some providers)"
+    return ""
+
+
+def _has_finetune_support(providers: List[KilnModelProvider]) -> str:
+    """Check if any provider supports fine-tuning"""
+    return "✅︎" if any(p.provider_finetune_id for p in providers) else ""
+
+
+@pytest.mark.paid(reason="Marking as paid so it isn't run by default")
+def test_generate_model_table():
+    """Generate a markdown table of all models and their capabilities"""
+
+    # Table header
+    table = [
+        "| Model Name | Providers | Structured Output | Reasoning | Synthetic Data | API Fine-Tuneable |",
+        "|------------|-----------|-------------------|-----------|----------------|-------------------|",
+    ]
+
+    for model in built_in_models:
+        provider_names = ", ".join(
+            sorted(provider_name_from_id(p.name.value) for p in model.providers)
+        )
+        structured_output = _get_support_status(
+            model.providers, "supports_structured_output"
+        )
+        reasoning = _get_support_status(model.providers, "reasoning_capable")
+        data_gen = _get_support_status(model.providers, "supports_data_gen")
+        finetune = _has_finetune_support(model.providers)
+
+        row = f"| {model.friendly_name} | {provider_names} | {structured_output} | {reasoning} | {data_gen} | {finetune} |"
+        table.append(row)
+
+    # Print the table (useful for documentation)
+    print("\nModel Capability Matrix:\n")
+    print("\n".join(table))
+
+    # Basic assertions to ensure the table is well-formed
+    assert len(table) > 2, "Table should have header and at least one row"
+    assert all("|" in row for row in table), "All rows should be properly formatted"
+    assert len(table[0].split("|")) == len(table[1].split("|")), (
+        "Header and separator should have same number of columns"
+    )

kiln_ai/adapters/test_prompt_adaptors.py

@@ -6,8 +6,8 @@ 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,
@@ -108,7 +108,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 +120,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,7 +131,7 @@ 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",
+        "model_provider": "ollama",
         "prompt_builder_name": "simple_prompt_builder",
     }
 

kiln_ai/adapters/test_prompt_builders.py

@@ -2,24 +2,31 @@ import json
 
 import pytest
 
-from kiln_ai.adapters.base_adapter import AdapterInfo, BaseAdapter
+from kiln_ai.adapters.model_adapters.base_adapter import AdapterInfo, 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,
     chain_of_thought_prompt,
     prompt_builder_from_ui_name,
 )
 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,
@@ -31,7 +38,7 @@ 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
@@ -62,7 +69,7 @@ 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 +77,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 +213,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 +254,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"},
@@ -260,7 +275,7 @@ def test_few_shot_prompt_builder(tmp_path):
                 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,7 +287,7 @@ 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)
@@ -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:
@@ -305,26 +320,84 @@ def test_prompt_builder_name():
     assert RepairsPromptBuilder.prompt_builder_name() == "repairs_prompt_builder"
 
 
-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
+def test_prompt_builder_from_ui_name(task_with_examples):
+    task = task_with_examples
+    assert isinstance(prompt_builder_from_ui_name("basic", task), SimplePromptBuilder)
+    assert isinstance(
+        prompt_builder_from_ui_name("few_shot", task), FewShotPromptBuilder
     )
-    assert (
-        prompt_builder_from_ui_name("few_shot_chain_of_thought")
-        == FewShotChainOfThoughtPromptBuilder
+    assert isinstance(
+        prompt_builder_from_ui_name("many_shot", task), MultiShotPromptBuilder
     )
-    assert (
-        prompt_builder_from_ui_name("multi_shot_chain_of_thought")
-        == MultiShotChainOfThoughtPromptBuilder
+    assert isinstance(
+        prompt_builder_from_ui_name("repairs", task), RepairsPromptBuilder
+    )
+    assert isinstance(
+        prompt_builder_from_ui_name("simple_chain_of_thought", task),
+        SimpleChainOfThoughtPromptBuilder,
+    )
+    assert isinstance(
+        prompt_builder_from_ui_name("few_shot_chain_of_thought", task),
+        FewShotChainOfThoughtPromptBuilder,
+    )
+    assert isinstance(
+        prompt_builder_from_ui_name("multi_shot_chain_of_thought", task),
+        MultiShotChainOfThoughtPromptBuilder,
     )
 
     with pytest.raises(ValueError, match="Unknown prompt builder: invalid_name"):
-        prompt_builder_from_ui_name("invalid_name")
+        prompt_builder_from_ui_name("invalid_name", task)
+
+    with pytest.raises(ValueError, match="Prompt ID not found: 123"):
+        prompt_builder_from_ui_name("id::123", task)
+
+    with pytest.raises(
+        ValueError,
+        match="Invalid fine-tune ID format. Expected 'project_id::task_id::fine_tune_id'",
+    ):
+        prompt_builder_from_ui_name("fine_tune_prompt::123", task)
+
+    with pytest.raises(
+        ValueError,
+        match="Fine-tune ID not found",
+    ):
+        prompt_builder_from_ui_name("fine_tune_prompt::123::456::789", task)
+
+    prompt = Prompt(
+        name="test_prompt_name",
+        prompt="test_prompt",
+        chain_of_thought_instructions="coti",
+        parent=task,
+    )
+    prompt.save_to_file()
+    pb = prompt_builder_from_ui_name("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,
+    )
+    finetune.save_to_file()
+    nested_fine_tune_id = (
+        task_with_examples.parent.id + "::" + task_with_examples.id + "::" + finetune.id
+    )
+    pb = prompt_builder_from_ui_name(
+        "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 +408,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 +476,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 +484,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 +496,94 @@ 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