kiln-ai 0.0.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

kiln_ai/adapters/repair/test_repair_task.py

@@ -0,0 +1,248 @@
+import json
+import os
+from unittest.mock import AsyncMock, patch
+
+import pytest
+from kiln_ai.adapters.langchain_adapters import (
+    LangChainPromptAdapter,
+)
+from kiln_ai.adapters.repair.repair_task import (
+    RepairTaskInput,
+    RepairTaskRun,
+)
+from kiln_ai.datamodel import (
+    DataSource,
+    DataSourceType,
+    Priority,
+    Task,
+    TaskOutput,
+    TaskRequirement,
+    TaskRun,
+)
+from pydantic import ValidationError
+
+json_joke_schema = """{
+  "type": "object",
+  "properties": {
+    "setup": {
+      "description": "The setup of the joke",
+      "title": "Setup",
+      "type": "string"
+    },
+    "punchline": {
+      "description": "The punchline to the joke",
+      "title": "Punchline",
+      "type": "string"
+    },
+    "rating": {
+      "anyOf": [
+        {
+          "type": "integer"
+        },
+        {
+          "type": "null"
+        }
+      ],
+      "default": null,
+      "description": "How funny the joke is, from 1 to 10",
+      "title": "Rating"
+    }
+  },
+  "required": [
+    "setup",
+    "punchline"
+  ]
+}
+"""
+
+
+@pytest.fixture
+def sample_task(tmp_path):
+    task_path = tmp_path / "task.json"
+    task = Task(
+        name="Joke Generator",
+        path=task_path,
+        description="Generate a funny joke",
+        instruction="Create a joke with a setup and punchline",
+        requirements=[
+            TaskRequirement(
+                id="req1",
+                name="Humor",
+                instruction="The joke should be funny and appropriate",
+                priority=Priority.p1,
+            )
+        ],
+        output_json_schema=json_joke_schema,
+    )
+    task.save_to_file()
+    return task
+
+
+@pytest.fixture
+def sample_task_run(sample_task):
+    task_run = TaskRun(
+        parent=sample_task,
+        input='{"topic": "chicken"}',
+        input_source=DataSource(
+            type=DataSourceType.human, properties={"created_by": "Jane Doe"}
+        ),
+        output=TaskOutput(
+            output='{"setup": "Why did the chicken cross the road?", "punchline": "To get to the other side", "rating": null}',
+            source=DataSource(
+                type=DataSourceType.synthetic,
+                properties={
+                    "model_name": "gpt_4o",
+                    "model_provider": "openai",
+                    "adapter_name": "langchain_adapter",
+                    "prompt_builder_name": "simple_prompt_builder",
+                },
+            ),
+        ),
+    )
+    task_run.save_to_file()
+    return task_run
+
+
+@pytest.fixture
+def sample_repair_data(sample_task, sample_task_run):
+    return {
+        "original_task": sample_task,
+        "task_run": sample_task_run,
+        "evaluator_feedback": "The joke is too cliché. Please come up with a more original chicken-related joke.",
+    }
+
+
+def test_build_repair_task_input(sample_repair_data):
+    result = RepairTaskRun.build_repair_task_input(**sample_repair_data)
+
+    assert isinstance(result, RepairTaskInput)
+    assert "Create a joke with a setup and punchline" in result.original_prompt
+    assert "1) The joke should be funny and appropriate" in result.original_prompt
+    assert result.original_input == '{"topic": "chicken"}'
+    assert (
+        result.original_output
+        == '{"setup": "Why did the chicken cross the road?", "punchline": "To get to the other side", "rating": null}'
+    )
+    assert (
+        result.evaluator_feedback
+        == "The joke is too cliché. Please come up with a more original chicken-related joke."
+    )
+
+
+def test_repair_input_schema():
+    schema = RepairTaskInput.model_json_schema()
+    assert schema["type"] == "object"
+    assert "original_prompt" in schema["properties"]
+    assert "original_input" in schema["properties"]
+    assert "original_output" in schema["properties"]
+    assert "evaluator_feedback" in schema["properties"]
+
+
+def test_repair_task_initialization(sample_task):
+    repair_task = RepairTaskRun(sample_task)
+
+    assert repair_task.name == "Repair"
+    assert "Repair a task run" in repair_task.description
+    assert "You are an assistant which helps improve output" in repair_task.instruction
+    assert len(repair_task.requirements) == 1
+    assert repair_task.requirements[0].name == "Follow Eval Feedback"
+    assert repair_task.input_json_schema == json.dumps(
+        RepairTaskInput.model_json_schema()
+    )
+    assert repair_task.output_json_schema == sample_task.output_json_schema
+
+
+def test_build_repair_task_input_with_empty_values(sample_task, sample_task_run):
+    # Arrange
+    sample_task_run.input = ""
+    sample_task_run.output.output = ""
+
+    # Act & Assert
+    with pytest.raises(ValidationError, match="evaluator_feedback"):
+        RepairTaskRun.build_repair_task_input(
+            original_task=sample_task, task_run=sample_task_run, evaluator_feedback=""
+        )
+
+    # Test that it works with non-empty feedback
+    result = RepairTaskRun.build_repair_task_input(
+        original_task=sample_task,
+        task_run=sample_task_run,
+        evaluator_feedback="Some feedback",
+    )
+    assert isinstance(result, RepairTaskInput)
+    assert result.evaluator_feedback == "Some feedback"
+
+
+@pytest.mark.parametrize("invalid_input", [{}])
+def test_build_repair_task_input_with_invalid_input(invalid_input):
+    # Act & Assert
+    with pytest.raises(TypeError):
+        RepairTaskRun.build_repair_task_input(invalid_input)
+
+
+@pytest.mark.paid
+async def test_live_run(sample_task, sample_task_run, sample_repair_data):
+    if os.getenv("GROQ_API_KEY") is None:
+        pytest.skip("GROQ_API_KEY not set")
+    repair_task = RepairTaskRun(sample_task)
+    repair_task_input = RepairTaskRun.build_repair_task_input(**sample_repair_data)
+    assert isinstance(repair_task_input, RepairTaskInput)
+
+    adapter = LangChainPromptAdapter(
+        repair_task, model_name="llama_3_1_8b", provider="groq"
+    )
+
+    run = await adapter.invoke(repair_task_input.model_dump())
+    assert run is not None
+    assert "Please come up with a more original chicken-related joke." in run.input
+    parsed_output = json.loads(run.output.output)
+    assert "setup" in parsed_output
+    assert "punchline" in parsed_output
+    assert run.output.source.properties == {
+        "adapter_name": "kiln_langchain_adapter",
+        "model_name": "llama_3_1_8b",
+        "model_provider": "groq",
+        "prompt_builder_name": "simple_prompt_builder",
+    }
+
+
+@pytest.mark.asyncio
+async def test_mocked_repair_task_run(sample_task, sample_task_run, sample_repair_data):
+    repair_task = RepairTaskRun(sample_task)
+    repair_task_input = RepairTaskRun.build_repair_task_input(**sample_repair_data)
+    assert isinstance(repair_task_input, RepairTaskInput)
+
+    mocked_output = {
+        "setup": "Why did the chicken join a band?",
+        "punchline": "Because it had excellent drumsticks!",
+        "rating": 8,
+    }
+
+    with patch.object(
+        LangChainPromptAdapter, "_run", new_callable=AsyncMock
+    ) as mock_run:
+        mock_run.return_value = mocked_output
+
+        adapter = LangChainPromptAdapter(
+            repair_task, model_name="llama_3_1_8b", provider="groq"
+        )
+
+        run = await adapter.invoke(repair_task_input.model_dump())
+
+    assert run is not None
+    assert run.id is None
+    assert "Please come up with a more original chicken-related joke." in run.input
+
+    parsed_output = json.loads(run.output.output)
+    assert parsed_output == mocked_output
+    assert run.output.source.properties == {
+        "adapter_name": "kiln_langchain_adapter",
+        "model_name": "llama_3_1_8b",
+        "model_provider": "groq",
+        "prompt_builder_name": "simple_prompt_builder",
+    }
+    assert run.input_source.type == DataSourceType.human
+    assert "created_by" in run.input_source.properties
+
+    # Verify that the mock was called
+    mock_run.assert_called_once()

kiln_ai/adapters/test_langchain_adapter.py

@@ -0,0 +1,50 @@
+from kiln_ai.adapters.langchain_adapters import LangChainPromptAdapter
+from kiln_ai.adapters.test_prompt_adaptors import build_test_task
+from langchain_groq import ChatGroq
+
+
+def test_langchain_adapter_munge_response(tmp_path):
+    task = build_test_task(tmp_path)
+    lca = LangChainPromptAdapter(
+        kiln_task=task, model_name="llama_3_1_8b", provider="ollama"
+    )
+    # Mistral Large tool calling format is a bit different
+    response = {
+        "name": "task_response",
+        "arguments": {
+            "setup": "Why did the cow join a band?",
+            "punchline": "Because she wanted to be a moo-sician!",
+        },
+    }
+    munged = lca._munge_response(response)
+    assert munged["setup"] == "Why did the cow join a band?"
+    assert munged["punchline"] == "Because she wanted to be a moo-sician!"
+
+    # non mistral format should continue to work
+    munged = lca._munge_response(response["arguments"])
+    assert munged["setup"] == "Why did the cow join a band?"
+    assert munged["punchline"] == "Because she wanted to be a moo-sician!"
+
+
+def test_langchain_adapter_infer_model_name(tmp_path):
+    task = build_test_task(tmp_path)
+    custom = ChatGroq(model="llama-3.1-8b-instant", groq_api_key="test")
+
+    lca = LangChainPromptAdapter(kiln_task=task, custom_model=custom)
+
+    model_info = lca.adapter_info()
+    assert model_info.model_name == "custom.langchain:llama-3.1-8b-instant"
+    assert model_info.model_provider == "custom.langchain:ChatGroq"
+
+
+def test_langchain_adapter_info(tmp_path):
+    task = build_test_task(tmp_path)
+
+    lca = LangChainPromptAdapter(
+        kiln_task=task, model_name="llama_3_1_8b", provider="ollama"
+    )
+
+    model_info = lca.adapter_info()
+    assert model_info.adapter_name == "kiln_langchain_adapter"
+    assert model_info.model_name == "llama_3_1_8b"
+    assert model_info.model_provider == "ollama"

kiln_ai/adapters/test_ml_model_list.py

@@ -0,0 +1,99 @@
+from unittest.mock import patch
+
+import pytest
+
+from libs.core.kiln_ai.adapters.ml_model_list import (
+    ModelProviderName,
+    check_provider_warnings,
+    provider_name_from_id,
+    provider_warnings,
+)
+
+
+@pytest.fixture
+def mock_config():
+    with patch("libs.core.kiln_ai.adapters.ml_model_list.get_config_value") as mock:
+        yield mock
+
+
+def test_check_provider_warnings_no_warning(mock_config):
+    mock_config.return_value = "some_value"
+
+    # This should not raise an exception
+    check_provider_warnings(ModelProviderName.amazon_bedrock)
+
+
+def test_check_provider_warnings_missing_key(mock_config):
+    mock_config.return_value = None
+
+    with pytest.raises(ValueError) as exc_info:
+        check_provider_warnings(ModelProviderName.amazon_bedrock)
+
+    assert provider_warnings[ModelProviderName.amazon_bedrock].message in str(
+        exc_info.value
+    )
+
+
+def test_check_provider_warnings_unknown_provider():
+    # This should not raise an exception, as no settings are required for unknown providers
+    check_provider_warnings("unknown_provider")
+
+
+@pytest.mark.parametrize(
+    "provider_name",
+    [
+        ModelProviderName.amazon_bedrock,
+        ModelProviderName.openrouter,
+        ModelProviderName.groq,
+        ModelProviderName.openai,
+    ],
+)
+def test_check_provider_warnings_all_providers(mock_config, provider_name):
+    mock_config.return_value = None
+
+    with pytest.raises(ValueError) as exc_info:
+        check_provider_warnings(provider_name)
+
+    assert provider_warnings[provider_name].message in str(exc_info.value)
+
+
+def test_check_provider_warnings_partial_keys_set(mock_config):
+    def mock_get(key):
+        return "value" if key == "bedrock_access_key" else None
+
+    mock_config.side_effect = mock_get
+
+    with pytest.raises(ValueError) as exc_info:
+        check_provider_warnings(ModelProviderName.amazon_bedrock)
+
+    assert provider_warnings[ModelProviderName.amazon_bedrock].message in str(
+        exc_info.value
+    )
+
+
+def test_provider_name_from_id_unknown_provider():
+    assert (
+        provider_name_from_id("unknown_provider")
+        == "Unknown provider: unknown_provider"
+    )
+
+
+def test_provider_name_from_id_case_sensitivity():
+    assert (
+        provider_name_from_id(ModelProviderName.amazon_bedrock.upper())
+        == "Unknown provider: AMAZON_BEDROCK"
+    )
+
+
+@pytest.mark.parametrize(
+    "provider_id, expected_name",
+    [
+        (ModelProviderName.amazon_bedrock, "Amazon Bedrock"),
+        (ModelProviderName.openrouter, "OpenRouter"),
+        (ModelProviderName.groq, "Groq"),
+        (ModelProviderName.ollama, "Ollama"),
+        (ModelProviderName.openai, "OpenAI"),
+    ],
+)
+def test_provider_name_from_id_parametrized(provider_id, expected_name):
+    assert provider_name_from_id(provider_id) == expected_name

kiln_ai/adapters/test_prompt_adaptors.py

@@ -0,0 +1,167 @@
+import os
+from pathlib import Path
+
+import kiln_ai.datamodel as datamodel
+import pytest
+from kiln_ai.adapters.langchain_adapters import LangChainPromptAdapter
+from kiln_ai.adapters.ml_model_list import built_in_models, ollama_online
+from langchain_core.language_models.fake_chat_models import FakeListChatModel
+
+
+@pytest.mark.paid
+async def test_groq(tmp_path):
+    if os.getenv("GROQ_API_KEY") is None:
+        pytest.skip("GROQ_API_KEY not set")
+    await run_simple_test(tmp_path, "llama_3_1_8b", "groq")
+
+
+@pytest.mark.paid
+async def test_openrouter(tmp_path):
+    await run_simple_test(tmp_path, "llama_3_1_8b", "openrouter")
+
+
+@pytest.mark.ollama
+async def test_ollama_phi(tmp_path):
+    # Check if Ollama API is running
+    if not await ollama_online():
+        pytest.skip("Ollama API not running. Expect it running on localhost:11434")
+
+    await run_simple_test(tmp_path, "phi_3_5", "ollama")
+
+
+@pytest.mark.ollama
+async def test_ollama_gemma(tmp_path):
+    # Check if Ollama API is running
+    if not await ollama_online():
+        pytest.skip("Ollama API not running. Expect it running on localhost:11434")
+
+    await run_simple_test(tmp_path, "gemma_2_2b", "ollama")
+
+
+@pytest.mark.ollama
+async def test_autoselect_provider(tmp_path):
+    # Check if Ollama API is running
+    if not await ollama_online():
+        pytest.skip("Ollama API not running. Expect it running on localhost:11434")
+
+    await run_simple_test(tmp_path, "phi_3_5")
+
+
+@pytest.mark.ollama
+async def test_ollama_llama(tmp_path):
+    # Check if Ollama API is running
+    if not await ollama_online():
+        pytest.skip("Ollama API not running. Expect it running on localhost:11434")
+
+    await run_simple_test(tmp_path, "llama_3_1_8b", "ollama")
+
+
+@pytest.mark.paid
+async def test_openai(tmp_path):
+    if os.getenv("OPENAI_API_KEY") is None:
+        pytest.skip("OPENAI_API_KEY not set")
+    await run_simple_test(tmp_path, "gpt_4o_mini", "openai")
+
+
+@pytest.mark.paid
+async def test_amazon_bedrock(tmp_path):
+    if (
+        os.getenv("AWS_SECRET_ACCESS_KEY") is None
+        or os.getenv("AWS_ACCESS_KEY_ID") is None
+    ):
+        pytest.skip("AWS keys not set")
+    await run_simple_test(tmp_path, "llama_3_1_8b", "amazon_bedrock")
+
+
+async def test_mock(tmp_path):
+    task = build_test_task(tmp_path)
+    mockChatModel = FakeListChatModel(responses=["mock response"])
+    adapter = LangChainPromptAdapter(task, custom_model=mockChatModel)
+    run = await adapter.invoke("You are a mock, send me the response!")
+    assert "mock response" in run.output.output
+
+
+async def test_mock_returning_run(tmp_path):
+    task = build_test_task(tmp_path)
+    mockChatModel = FakeListChatModel(responses=["mock response"])
+    adapter = LangChainPromptAdapter(task, custom_model=mockChatModel)
+    run = await adapter.invoke("You are a mock, send me the response!")
+    assert run.output.output == "mock response"
+    assert run is not None
+    assert run.id is not None
+    assert run.input == "You are a mock, send me the response!"
+    assert run.output.output == "mock response"
+    assert "created_by" in run.input_source.properties
+    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",
+    }
+
+
+@pytest.mark.paid
+@pytest.mark.ollama
+async def test_all_built_in_models(tmp_path):
+    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
+
+
+def build_test_task(tmp_path: Path):
+    project = datamodel.Project(name="test", path=tmp_path / "test.kiln")
+    project.save_to_file()
+    assert project.name == "test"
+
+    r1 = datamodel.TaskRequirement(
+        name="BEDMAS",
+        instruction="You follow order of mathematical operation (BEDMAS)",
+    )
+    r2 = datamodel.TaskRequirement(
+        name="only basic math",
+        instruction="If the problem has anything other than addition, subtraction, multiplication, division, and brackets, you will not answer it. Reply instead with 'I'm just a basic calculator, I don't know how to do that'.",
+    )
+    r3 = datamodel.TaskRequirement(
+        name="Answer format",
+        instruction="The answer can contain any content about your reasoning, but at the end it should include the final answer in numerals in square brackets. For example if the answer is one hundred, the end of your response should be [100].",
+    )
+    task = datamodel.Task(
+        parent=project,
+        name="test task",
+        instruction="You are an assistant which performs math tasks provided in plain text.",
+        requirements=[r1, r2, r3],
+    )
+    task.save_to_file()
+    assert task.name == "test task"
+    assert len(task.requirements) == 3
+    return task
+
+
+async def run_simple_test(tmp_path: Path, model_name: str, provider: str | None = None):
+    task = build_test_task(tmp_path)
+    return await run_simple_task(task, model_name, provider)
+
+
+async def run_simple_task(task: datamodel.Task, model_name: str, provider: str):
+    adapter = LangChainPromptAdapter(task, model_name=model_name, provider=provider)
+
+    run = await adapter.invoke(
+        "You should answer the following question: four plus six times 10"
+    )
+    assert "64" in run.output.output
+    assert run.id is not None
+    assert (
+        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",
+    }