kiln-ai 0.5.4__py3-none-any.whl → 0.6.0__py3-none-any.whl

kiln_ai/adapters/ml_model_list.py

@@ -69,6 +69,7 @@ class ModelName(str, Enum):
     gemma_2_2b = "gemma_2_2b"
     gemma_2_9b = "gemma_2_9b"
     gemma_2_27b = "gemma_2_27b"
+    claude_3_5_haiku = "claude_3_5_haiku"
     claude_3_5_sonnet = "claude_3_5_sonnet"
     gemini_1_5_flash = "gemini_1_5_flash"
     gemini_1_5_flash_8b = "gemini_1_5_flash_8b"
@@ -88,6 +89,7 @@ class KilnModelProvider(BaseModel):
 
     name: ModelProviderName
     supports_structured_output: bool = True
+    supports_data_gen: bool = True
     provider_options: Dict = {}
 
 
@@ -143,6 +145,18 @@ built_in_models: List[KilnModel] = [
             ),
         ],
     ),
+    # Claude 3.5 Haiku
+    KilnModel(
+        family=ModelFamily.claude,
+        name=ModelName.claude_3_5_haiku,
+        friendly_name="Claude 3.5 Haiku",
+        providers=[
+            KilnModelProvider(
+                name=ModelProviderName.openrouter,
+                provider_options={"model": "anthropic/claude-3-5-haiku"},
+            ),
+        ],
+    ),
     # Claude 3.5 Sonnet
     KilnModel(
         family=ModelFamily.claude,
@@ -163,6 +177,8 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_structured_output=False,  # it should, but doesn't work on openrouter
+                supports_data_gen=False,  # doesn't work on openrouter
                 provider_options={"model": "google/gemini-pro-1.5"},
             ),
         ],
@@ -175,6 +191,7 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_data_gen=False,
                 provider_options={"model": "google/gemini-flash-1.5"},
             ),
         ],
@@ -187,6 +204,8 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "google/gemini-flash-1.5-8b"},
             ),
         ],
@@ -200,6 +219,7 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "nvidia/llama-3.1-nemotron-70b-instruct"},
             ),
         ],
@@ -217,6 +237,7 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.amazon_bedrock,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={
                     "model": "meta.llama3-1-8b-instruct-v1:0",
                     "region_name": "us-west-2",  # Llama 3.1 only in west-2
@@ -224,6 +245,7 @@ built_in_models: List[KilnModel] = [
             ),
             KilnModelProvider(
                 name=ModelProviderName.ollama,
+                supports_data_gen=False,
                 provider_options={
                     "model": "llama3.1:8b",
                     "model_aliases": ["llama3.1"],  # 8b is default
@@ -232,6 +254,7 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "meta-llama/llama-3.1-8b-instruct"},
             ),
         ],
@@ -248,7 +271,9 @@ built_in_models: List[KilnModel] = [
             ),
             KilnModelProvider(
                 name=ModelProviderName.amazon_bedrock,
+                # not sure how AWS manages to break this, but it's not working
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={
                     "model": "meta.llama3-1-70b-instruct-v1:0",
                     "region_name": "us-west-2",  # Llama 3.1 only in west-2
@@ -272,6 +297,7 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.amazon_bedrock,
+                supports_data_gen=False,
                 provider_options={
                     "model": "meta.llama3-1-405b-instruct-v1:0",
                     "region_name": "us-west-2",  # Llama 3.1 only in west-2
@@ -331,8 +357,15 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "meta-llama/llama-3.2-3b-instruct"},
             ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                supports_structured_output=False,
+                supports_data_gen=False,
+                provider_options={"model": "llama3.2"},
+            ),
         ],
     ),
     # Llama 3.2 11B
@@ -344,8 +377,15 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "meta-llama/llama-3.2-11b-vision-instruct"},
             ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                supports_structured_output=False,
+                supports_data_gen=False,
+                provider_options={"model": "llama3.2-vision"},
+            ),
         ],
     ),
     # Llama 3.2 90B
@@ -357,8 +397,15 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "meta-llama/llama-3.2-90b-vision-instruct"},
             ),
+            KilnModelProvider(
+                name=ModelProviderName.ollama,
+                supports_structured_output=False,
+                supports_data_gen=False,
+                provider_options={"model": "llama3.2-vision:90b"},
+            ),
         ],
     ),
     # Phi 3.5
@@ -371,10 +418,13 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.ollama,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "phi3.5"},
             ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={"model": "microsoft/phi-3.5-mini-128k-instruct"},
             ),
         ],
@@ -389,6 +439,7 @@ built_in_models: List[KilnModel] = [
             KilnModelProvider(
                 name=ModelProviderName.ollama,
                 supports_structured_output=False,
+                supports_data_gen=False,
                 provider_options={
                     "model": "gemma2:2b",
                 },
@@ -404,12 +455,14 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.ollama,
+                supports_data_gen=False,
                 provider_options={
                     "model": "gemma2:9b",
                 },
             ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_data_gen=False,
                 provider_options={"model": "google/gemma-2-9b-it"},
             ),
         ],
@@ -423,12 +476,14 @@ built_in_models: List[KilnModel] = [
         providers=[
             KilnModelProvider(
                 name=ModelProviderName.ollama,
+                supports_data_gen=False,
                 provider_options={
                     "model": "gemma2:27b",
                 },
             ),
             KilnModelProvider(
                 name=ModelProviderName.openrouter,
+                supports_data_gen=False,
                 provider_options={"model": "google/gemma-2-27b-it"},
             ),
         ],
@@ -436,6 +491,19 @@ built_in_models: List[KilnModel] = [
 ]
 
 
+def get_model_and_provider(
+    model_name: str, provider_name: str
+) -> tuple[KilnModel | None, KilnModelProvider | None]:
+    model = next(filter(lambda m: m.name == model_name, built_in_models), None)
+    if model is None:
+        return None, None
+    provider = next(filter(lambda p: p.name == provider_name, model.providers), None)
+    # all or nothing
+    if provider is None or model is None:
+        return None, None
+    return model, provider
+
+
 def provider_name_from_id(id: str) -> str:
     """
     Converts a provider ID to its human-readable name.
@@ -674,7 +742,6 @@ def parse_ollama_tags(tags: Any) -> OllamaConnection | None:
         models = tags["models"]
         if isinstance(models, list):
             model_names = [model["model"] for model in models]
-            print(f"model_names: {model_names}")
             available_supported_models = [
                 model
                 for model in model_names

kiln_ai/adapters/prompt_builders.py

@@ -54,6 +54,28 @@ class BasePromptBuilder(metaclass=ABCMeta):
 
         return f"The input is:\n{input}"
 
+    def chain_of_thought_prompt(self) -> str | None:
+        """Build and return the chain of thought prompt string.
+
+        Returns:
+            str: The constructed chain of thought prompt.
+        """
+        return None
+
+    def build_prompt_for_ui(self) -> str:
+        """Build a prompt for the UI. It includes additional instructions (like chain of thought), even if they are passed to the model in stages.
+
+        Designed for end-user consumption, not for model consumption.
+
+        Returns:
+            str: The constructed prompt string.
+        """
+        base_prompt = self.build_prompt()
+        cot_prompt = self.chain_of_thought_prompt()
+        if cot_prompt:
+            base_prompt += "\n# Thinking Instructions\n\n" + cot_prompt
+        return base_prompt
+
 
 class SimplePromptBuilder(BasePromptBuilder):
     """A basic prompt builder that combines task instruction with requirements."""
@@ -187,11 +209,49 @@ class RepairsPromptBuilder(MultiShotPromptBuilder):
         return prompt_section
 
 
+def chain_of_thought_prompt(task: Task) -> str | None:
+    """Standard implementation to build and return the chain of thought prompt string.
+
+    Returns:
+        str: The constructed chain of thought prompt.
+    """
+
+    cot_instruction = task.thinking_instruction
+    if not cot_instruction:
+        cot_instruction = "Think step by step, explaining your reasoning."
+
+    return cot_instruction
+
+
+class SimpleChainOfThoughtPromptBuilder(SimplePromptBuilder):
+    """A prompt builder that includes a chain of thought prompt on top of the simple prompt."""
+
+    def chain_of_thought_prompt(self) -> str | None:
+        return chain_of_thought_prompt(self.task)
+
+
+class FewShotChainOfThoughtPromptBuilder(FewShotPromptBuilder):
+    """A prompt builder that includes a chain of thought prompt on top of the few shot prompt."""
+
+    def chain_of_thought_prompt(self) -> str | None:
+        return chain_of_thought_prompt(self.task)
+
+
+class MultiShotChainOfThoughtPromptBuilder(MultiShotPromptBuilder):
+    """A prompt builder that includes a chain of thought prompt on top of the multi shot prompt."""
+
+    def chain_of_thought_prompt(self) -> str | None:
+        return chain_of_thought_prompt(self.task)
+
+
 prompt_builder_registry = {
     "simple_prompt_builder": SimplePromptBuilder,
     "multi_shot_prompt_builder": MultiShotPromptBuilder,
     "few_shot_prompt_builder": FewShotPromptBuilder,
     "repairs_prompt_builder": RepairsPromptBuilder,
+    "simple_chain_of_thought_prompt_builder": SimpleChainOfThoughtPromptBuilder,
+    "few_shot_chain_of_thought_prompt_builder": FewShotChainOfThoughtPromptBuilder,
+    "multi_shot_chain_of_thought_prompt_builder": MultiShotChainOfThoughtPromptBuilder,
 }
 
 
@@ -217,5 +277,11 @@ def prompt_builder_from_ui_name(ui_name: str) -> type[BasePromptBuilder]:
             return MultiShotPromptBuilder
         case "repairs":
             return RepairsPromptBuilder
+        case "simple_chain_of_thought":
+            return SimpleChainOfThoughtPromptBuilder
+        case "few_shot_chain_of_thought":
+            return FewShotChainOfThoughtPromptBuilder
+        case "multi_shot_chain_of_thought":
+            return MultiShotChainOfThoughtPromptBuilder
         case _:
             raise ValueError(f"Unknown prompt builder: {ui_name}")

kiln_ai/adapters/repair/test_repair_task.py

@@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, patch
 import pytest
 from pydantic import ValidationError
 
+from kiln_ai.adapters.base_adapter import RunOutput
 from kiln_ai.adapters.langchain_adapters import (
     LangChainPromptAdapter,
 )
@@ -222,7 +223,9 @@ async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repai
     with patch.object(
         LangChainPromptAdapter, "_run", new_callable=AsyncMock
     ) as mock_run:
-        mock_run.return_value = mocked_output
+        mock_run.return_value = RunOutput(
+            output=mocked_output, intermediate_outputs=None
+        )
 
         adapter = LangChainPromptAdapter(
             repair_task, model_name="llama_3_1_8b", provider="groq"

kiln_ai/adapters/test_langchain_adapter.py

@@ -1,6 +1,10 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 from langchain_groq import ChatGroq
 
 from kiln_ai.adapters.langchain_adapters import LangChainPromptAdapter
+from kiln_ai.adapters.prompt_builders import SimpleChainOfThoughtPromptBuilder
 from kiln_ai.adapters.test_prompt_adaptors import build_test_task
 
 
@@ -49,3 +53,72 @@ def test_langchain_adapter_info(tmp_path):
     assert model_info.adapter_name == "kiln_langchain_adapter"
     assert model_info.model_name == "llama_3_1_8b"
     assert model_info.model_provider == "ollama"
+
+
+async def test_langchain_adapter_with_cot(tmp_path):
+    task = build_test_task(tmp_path)
+    task.output_json_schema = (
+        '{"type": "object", "properties": {"count": {"type": "integer"}}}'
+    )
+    lca = LangChainPromptAdapter(
+        kiln_task=task,
+        model_name="llama_3_1_8b",
+        provider="ollama",
+        prompt_builder=SimpleChainOfThoughtPromptBuilder(task),
+    )
+
+    # Mock the base model and its invoke method
+    mock_base_model = MagicMock()
+    mock_base_model.invoke.return_value = AIMessage(
+        content="Chain of thought reasoning..."
+    )
+
+    # Create a separate mock for self.model()
+    mock_model_instance = MagicMock()
+    mock_model_instance.invoke.return_value = {"parsed": {"count": 1}}
+
+    # Mock the langchain_model_from function to return the base model
+    mock_model_from = AsyncMock(return_value=mock_base_model)
+
+    # Patch both the langchain_model_from function and self.model()
+    with (
+        patch(
+            "kiln_ai.adapters.langchain_adapters.langchain_model_from", mock_model_from
+        ),
+        patch.object(LangChainPromptAdapter, "model", return_value=mock_model_instance),
+    ):
+        response = await lca._run("test input")
+
+    # First 3 messages are the same for both calls
+    for invoke_args in [
+        mock_base_model.invoke.call_args[0][0],
+        mock_model_instance.invoke.call_args[0][0],
+    ]:
+        assert isinstance(
+            invoke_args[0], SystemMessage
+        )  # First message should be system prompt
+        assert (
+            "You are an assistant which performs math tasks provided in plain text."
+            in invoke_args[0].content
+        )
+        assert isinstance(invoke_args[1], HumanMessage)
+        assert "test input" in invoke_args[1].content
+        assert isinstance(invoke_args[2], SystemMessage)
+        assert "step by step" in invoke_args[2].content
+
+    # the COT should only have 3 messages
+    assert len(mock_base_model.invoke.call_args[0][0]) == 3
+    assert len(mock_model_instance.invoke.call_args[0][0]) == 5
+
+    # the final response should have the COT content and the final instructions
+    invoke_args = mock_model_instance.invoke.call_args[0][0]
+    assert isinstance(invoke_args[3], AIMessage)
+    assert "Chain of thought reasoning..." in invoke_args[3].content
+    assert isinstance(invoke_args[4], SystemMessage)
+    assert "Considering the above, return a final result." in invoke_args[4].content
+
+    assert (
+        response.intermediate_outputs["chain_of_thought"]
+        == "Chain of thought reasoning..."
+    )
+    assert response.output == {"count": 1}

kiln_ai/adapters/test_ml_model_list.py

@@ -4,9 +4,11 @@ from unittest.mock import patch
 import pytest
 
 from kiln_ai.adapters.ml_model_list import (
+    ModelName,
     ModelProviderName,
     OllamaConnection,
     check_provider_warnings,
+    get_model_and_provider,
     ollama_model_supported,
     parse_ollama_tags,
     provider_name_from_id,
@@ -123,3 +125,57 @@ def test_ollama_model_supported():
     assert ollama_model_supported(conn, "llama3.1:latest")
     assert ollama_model_supported(conn, "llama3.1")
     assert not ollama_model_supported(conn, "unknown_model")
+
+
+def test_get_model_and_provider_valid():
+    # Test with a known valid model and provider combination
+    model, provider = get_model_and_provider(
+        ModelName.phi_3_5, ModelProviderName.ollama
+    )
+
+    assert model is not None
+    assert provider is not None
+    assert model.name == ModelName.phi_3_5
+    assert provider.name == ModelProviderName.ollama
+    assert provider.provider_options["model"] == "phi3.5"
+
+
+def test_get_model_and_provider_invalid_model():
+    # Test with an invalid model name
+    model, provider = get_model_and_provider(
+        "nonexistent_model", ModelProviderName.ollama
+    )
+
+    assert model is None
+    assert provider is None
+
+
+def test_get_model_and_provider_invalid_provider():
+    # Test with a valid model but invalid provider
+    model, provider = get_model_and_provider(ModelName.phi_3_5, "nonexistent_provider")
+
+    assert model is None
+    assert provider is None
+
+
+def test_get_model_and_provider_valid_model_wrong_provider():
+    # Test with a valid model but a provider that doesn't support it
+    model, provider = get_model_and_provider(
+        ModelName.phi_3_5, ModelProviderName.amazon_bedrock
+    )
+
+    assert model is None
+    assert provider is None
+
+
+def test_get_model_and_provider_multiple_providers():
+    # Test with a model that has multiple providers
+    model, provider = get_model_and_provider(
+        ModelName.llama_3_1_70b, ModelProviderName.groq
+    )
+
+    assert model is not None
+    assert provider is not None
+    assert model.name == ModelName.llama_3_1_70b
+    assert provider.name == ModelProviderName.groq
+    assert provider.provider_options["model"] == "llama-3.1-70b-versatile"

kiln_ai/adapters/test_prompt_adaptors.py

@@ -7,6 +7,18 @@ from langchain_core.language_models.fake_chat_models import FakeListChatModel
 import kiln_ai.datamodel as datamodel
 from kiln_ai.adapters.langchain_adapters import LangChainPromptAdapter
 from kiln_ai.adapters.ml_model_list import built_in_models, ollama_online
+from kiln_ai.adapters.prompt_builders import (
+    BasePromptBuilder,
+    SimpleChainOfThoughtPromptBuilder,
+)
+
+
+def get_all_models_and_providers():
+    model_provider_pairs = []
+    for model in built_in_models:
+        for provider in model.providers:
+            model_provider_pairs.append((model.name, provider.name))
+    return model_provider_pairs
 
 
 @pytest.mark.paid
@@ -28,6 +40,9 @@ async def test_groq(tmp_path):
         "llama_3_2_3b",
         "llama_3_2_11b",
         "llama_3_2_90b",
+        "claude_3_5_haiku",
+        "claude_3_5_sonnet",
+        "phi_3_5",
     ],
 )
 @pytest.mark.paid
@@ -117,15 +132,19 @@ async def test_mock_returning_run(tmp_path):
 
 @pytest.mark.paid
 @pytest.mark.ollama
-async def test_all_built_in_models(tmp_path):
+@pytest.mark.parametrize("model_name,provider_name", get_all_models_and_providers())
+async def test_all_models_providers_plaintext(tmp_path, model_name, provider_name):
     task = build_test_task(tmp_path)
-    for model in built_in_models:
-        for provider in model.providers:
-            try:
-                print(f"Running {model.name} {provider.name}")
-                await run_simple_task(task, model.name, provider.name)
-            except Exception as e:
-                raise RuntimeError(f"Error running {model.name} {provider}") from e
+    await run_simple_task(task, model_name, provider_name)
+
+
+@pytest.mark.paid
+@pytest.mark.ollama
+@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)
 
 
 def build_test_task(tmp_path: Path):
@@ -157,13 +176,25 @@ def build_test_task(tmp_path: Path):
     return task
 
 
-async def run_simple_test(tmp_path: Path, model_name: str, provider: str | None = None):
+async def run_simple_test(
+    tmp_path: Path,
+    model_name: str,
+    provider: str | None = None,
+    prompt_builder: BasePromptBuilder | None = None,
+):
     task = build_test_task(tmp_path)
-    return await run_simple_task(task, model_name, provider)
+    return await run_simple_task(task, model_name, provider, prompt_builder)
 
 
-async def run_simple_task(task: datamodel.Task, model_name: str, provider: str):
-    adapter = LangChainPromptAdapter(task, model_name=model_name, provider=provider)
+async def run_simple_task(
+    task: datamodel.Task,
+    model_name: str,
+    provider: str,
+    prompt_builder: BasePromptBuilder | None = None,
+) -> datamodel.TaskRun:
+    adapter = LangChainPromptAdapter(
+        task, model_name=model_name, provider=provider, prompt_builder=prompt_builder
+    )
 
     run = await adapter.invoke(
         "You should answer the following question: four plus six times 10"
@@ -174,9 +205,14 @@ async def run_simple_task(task: datamodel.Task, model_name: str, provider: str):
         run.input == "You should answer the following question: four plus six times 10"
     )
     assert "64" in run.output.output
-    assert run.output.source.properties == {
-        "adapter_name": "kiln_langchain_adapter",
-        "model_name": model_name,
-        "model_provider": provider,
-        "prompt_builder_name": "simple_prompt_builder",
-    }
+    source_props = run.output.source.properties
+    assert source_props["adapter_name"] == "kiln_langchain_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
+    return run