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

kiln_ai/adapters/model_adapters/test_litellm_adapter.py

@@ -0,0 +1,407 @@
+import json
+from unittest.mock import Mock, patch
+
+import pytest
+
+from kiln_ai.adapters.ml_model_list import ModelProviderName, StructuredOutputMode
+from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig
+from kiln_ai.adapters.model_adapters.litellm_adapter import LiteLlmAdapter
+from kiln_ai.adapters.model_adapters.litellm_config import (
+    LiteLlmConfig,
+)
+from kiln_ai.datamodel import Project, Task
+
+
+@pytest.fixture
+def mock_task(tmp_path):
+    # Create a project first since Task requires a parent
+    project_path = tmp_path / "test_project" / "project.kiln"
+    project_path.parent.mkdir()
+
+    project = Project(name="Test Project", path=str(project_path))
+    project.save_to_file()
+
+    schema = {
+        "type": "object",
+        "properties": {"test": {"type": "string"}},
+    }
+
+    task = Task(
+        name="Test Task",
+        instruction="Test instruction",
+        parent=project,
+        output_json_schema=json.dumps(schema),
+    )
+    task.save_to_file()
+    return task
+
+
+@pytest.fixture
+def config():
+    return LiteLlmConfig(
+        base_url="https://api.test.com",
+        model_name="test-model",
+        provider_name="openrouter",
+        default_headers={"X-Test": "test"},
+        additional_body_options={"api_key": "test_key"},
+    )
+
+
+def test_initialization(config, mock_task):
+    adapter = LiteLlmAdapter(
+        config=config,
+        kiln_task=mock_task,
+        prompt_id="simple_prompt_builder",
+        base_adapter_config=AdapterConfig(default_tags=["test-tag"]),
+    )
+
+    assert adapter.config == config
+    assert adapter.run_config.task == mock_task
+    assert adapter.run_config.prompt_id == "simple_prompt_builder"
+    assert adapter.base_adapter_config.default_tags == ["test-tag"]
+    assert adapter.run_config.model_name == config.model_name
+    assert adapter.run_config.model_provider_name == config.provider_name
+    assert adapter.config.additional_body_options["api_key"] == "test_key"
+    assert adapter._api_base == config.base_url
+    assert adapter._headers == config.default_headers
+
+
+def test_adapter_info(config, mock_task):
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    assert adapter.adapter_name() == "kiln_openai_compatible_adapter"
+
+    assert adapter.run_config.model_name == config.model_name
+    assert adapter.run_config.model_provider_name == config.provider_name
+    assert adapter.run_config.prompt_id == "simple_prompt_builder"
+
+
+@pytest.mark.asyncio
+async def test_response_format_options_unstructured(config, mock_task):
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    # Mock has_structured_output to return False
+    with patch.object(adapter, "has_structured_output", return_value=False):
+        options = await adapter.response_format_options()
+        assert options == {}
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [
+        StructuredOutputMode.json_mode,
+        StructuredOutputMode.json_instruction_and_object,
+    ],
+)
+@pytest.mark.asyncio
+async def test_response_format_options_json_mode(config, mock_task, mode):
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    with (
+        patch.object(adapter, "has_structured_output", return_value=True),
+        patch.object(adapter, "model_provider") as mock_provider,
+    ):
+        mock_provider.return_value.structured_output_mode = mode
+
+        options = await adapter.response_format_options()
+        assert options == {"response_format": {"type": "json_object"}}
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [
+        StructuredOutputMode.default,
+        StructuredOutputMode.function_calling,
+    ],
+)
+@pytest.mark.asyncio
+async def test_response_format_options_function_calling(config, mock_task, mode):
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    with (
+        patch.object(adapter, "has_structured_output", return_value=True),
+        patch.object(adapter, "model_provider") as mock_provider,
+    ):
+        mock_provider.return_value.structured_output_mode = mode
+
+        options = await adapter.response_format_options()
+        assert "tools" in options
+        # full tool structure validated below
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [
+        StructuredOutputMode.json_custom_instructions,
+        StructuredOutputMode.json_instructions,
+    ],
+)
+@pytest.mark.asyncio
+async def test_response_format_options_json_instructions(config, mock_task, mode):
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    with (
+        patch.object(adapter, "has_structured_output", return_value=True),
+        patch.object(adapter, "model_provider") as mock_provider,
+    ):
+        mock_provider.return_value.structured_output_mode = (
+            StructuredOutputMode.json_instructions
+        )
+        options = await adapter.response_format_options()
+        assert options == {}
+
+
+@pytest.mark.asyncio
+async def test_response_format_options_json_schema(config, mock_task):
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    with (
+        patch.object(adapter, "has_structured_output", return_value=True),
+        patch.object(adapter, "model_provider") as mock_provider,
+    ):
+        mock_provider.return_value.structured_output_mode = (
+            StructuredOutputMode.json_schema
+        )
+        options = await adapter.response_format_options()
+        assert options == {
+            "response_format": {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "task_response",
+                    "schema": mock_task.output_schema(),
+                },
+            }
+        }
+
+
+def test_tool_call_params_weak(config, mock_task):
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    params = adapter.tool_call_params(strict=False)
+    expected_schema = mock_task.output_schema()
+    expected_schema["additionalProperties"] = False
+
+    assert params == {
+        "tools": [
+            {
+                "type": "function",
+                "function": {
+                    "name": "task_response",
+                    "parameters": expected_schema,
+                },
+            }
+        ],
+        "tool_choice": {
+            "type": "function",
+            "function": {"name": "task_response"},
+        },
+    }
+
+
+def test_tool_call_params_strict(config, mock_task):
+    config.provider_name = "openai"
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    params = adapter.tool_call_params(strict=True)
+    expected_schema = mock_task.output_schema()
+    expected_schema["additionalProperties"] = False
+
+    assert params == {
+        "tools": [
+            {
+                "type": "function",
+                "function": {
+                    "name": "task_response",
+                    "parameters": expected_schema,
+                    "strict": True,
+                },
+            }
+        ],
+        "tool_choice": {
+            "type": "function",
+            "function": {"name": "task_response"},
+        },
+    }
+
+
+@pytest.mark.parametrize(
+    "provider_name,expected_prefix",
+    [
+        (ModelProviderName.openrouter, "openrouter"),
+        (ModelProviderName.openai, "openai"),
+        (ModelProviderName.groq, "groq"),
+        (ModelProviderName.anthropic, "anthropic"),
+        (ModelProviderName.ollama, "openai"),
+        (ModelProviderName.gemini_api, "gemini"),
+        (ModelProviderName.fireworks_ai, "fireworks_ai"),
+        (ModelProviderName.amazon_bedrock, "bedrock"),
+        (ModelProviderName.azure_openai, "azure"),
+        (ModelProviderName.huggingface, "huggingface"),
+        (ModelProviderName.vertex, "vertex_ai"),
+        (ModelProviderName.together_ai, "together_ai"),
+    ],
+)
+def test_litellm_model_id_standard_providers(
+    config, mock_task, provider_name, expected_prefix
+):
+    """Test litellm_model_id for standard providers"""
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    # Mock the model_provider method to return a provider with the specified name
+    mock_provider = Mock()
+    mock_provider.name = provider_name
+    mock_provider.model_id = "test-model"
+
+    with patch.object(adapter, "model_provider", return_value=mock_provider):
+        model_id = adapter.litellm_model_id()
+
+    assert model_id == f"{expected_prefix}/test-model"
+    # Verify caching works
+    assert adapter._litellm_model_id == model_id
+
+
+@pytest.mark.parametrize(
+    "provider_name",
+    [
+        ModelProviderName.openai_compatible,
+        ModelProviderName.kiln_custom_registry,
+        ModelProviderName.kiln_fine_tune,
+    ],
+)
+def test_litellm_model_id_custom_providers(config, mock_task, provider_name):
+    """Test litellm_model_id for custom providers that require a base URL"""
+    config.base_url = "https://api.custom.com"
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    # Mock the model_provider method
+    mock_provider = Mock()
+    mock_provider.name = provider_name
+    mock_provider.model_id = "custom-model"
+
+    with patch.object(adapter, "model_provider", return_value=mock_provider):
+        model_id = adapter.litellm_model_id()
+
+    # Custom providers should use "openai" as the provider name
+    assert model_id == "openai/custom-model"
+    assert adapter._litellm_model_id == model_id
+
+
+def test_litellm_model_id_custom_provider_no_base_url(config, mock_task):
+    """Test litellm_model_id raises error for custom providers without base URL"""
+    config.base_url = None
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    # Mock the model_provider method
+    mock_provider = Mock()
+    mock_provider.name = ModelProviderName.openai_compatible
+    mock_provider.model_id = "custom-model"
+
+    with patch.object(adapter, "model_provider", return_value=mock_provider):
+        with pytest.raises(ValueError, match="Explicit Base URL is required"):
+            adapter.litellm_model_id()
+
+
+def test_litellm_model_id_no_model_id(config, mock_task):
+    """Test litellm_model_id raises error when provider has no model_id"""
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    # Mock the model_provider method to return a provider with no model_id
+    mock_provider = Mock()
+    mock_provider.name = ModelProviderName.openai
+    mock_provider.model_id = None
+
+    with patch.object(adapter, "model_provider", return_value=mock_provider):
+        with pytest.raises(ValueError, match="Model ID is required"):
+            adapter.litellm_model_id()
+
+
+def test_litellm_model_id_caching(config, mock_task):
+    """Test that litellm_model_id caches the result"""
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    # Set the cached value directly
+    adapter._litellm_model_id = "cached-value"
+
+    # The method should return the cached value without calling model_provider
+    with patch.object(adapter, "model_provider") as mock_model_provider:
+        model_id = adapter.litellm_model_id()
+
+    assert model_id == "cached-value"
+    mock_model_provider.assert_not_called()
+
+
+def test_litellm_model_id_unknown_provider(config, mock_task):
+    """Test litellm_model_id raises error for unknown provider"""
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+
+    # Create a mock provider with an unknown name
+    mock_provider = Mock()
+    mock_provider.name = "unknown_provider"  # Not in ModelProviderName enum
+    mock_provider.model_id = "test-model"
+
+    with patch.object(adapter, "model_provider", return_value=mock_provider):
+        with patch(
+            "kiln_ai.adapters.model_adapters.litellm_adapter.raise_exhaustive_enum_error"
+        ) as mock_raise_error:
+            mock_raise_error.side_effect = Exception("Test error")
+
+            with pytest.raises(Exception, match="Test error"):
+                adapter.litellm_model_id()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "top_logprobs,response_format,extra_body",
+    [
+        (None, {}, {}),  # Basic case
+        (5, {}, {}),  # With logprobs
+        (
+            None,
+            {"response_format": {"type": "json_object"}},
+            {},
+        ),  # With response format
+        (
+            3,
+            {"tools": [{"type": "function"}]},
+            {"reasoning_effort": 0.8},
+        ),  # Combined options
+    ],
+)
+async def test_build_completion_kwargs(
+    config, mock_task, top_logprobs, response_format, extra_body
+):
+    """Test build_completion_kwargs with various configurations"""
+    adapter = LiteLlmAdapter(config=config, kiln_task=mock_task)
+    mock_provider = Mock()
+    messages = [{"role": "user", "content": "Hello"}]
+
+    with (
+        patch.object(adapter, "model_provider", return_value=mock_provider),
+        patch.object(adapter, "litellm_model_id", return_value="openai/test-model"),
+        patch.object(adapter, "build_extra_body", return_value=extra_body),
+        patch.object(adapter, "response_format_options", return_value=response_format),
+    ):
+        kwargs = await adapter.build_completion_kwargs(
+            mock_provider, messages, top_logprobs
+        )
+
+    # Verify core functionality
+    assert kwargs["model"] == "openai/test-model"
+    assert kwargs["messages"] == messages
+    assert kwargs["api_base"] == config.base_url
+
+    # Verify optional parameters
+    if top_logprobs is not None:
+        assert kwargs["logprobs"] is True
+        assert kwargs["top_logprobs"] == top_logprobs
+    else:
+        assert "logprobs" not in kwargs
+        assert "top_logprobs" not in kwargs
+
+    # Verify response format is included
+    for key, value in response_format.items():
+        assert kwargs[key] == value
+
+    # Verify extra body is included
+    for key, value in extra_body.items():
+        assert kwargs[key] == value

kiln_ai/adapters/model_adapters/test_structured_output.py

@@ -66,7 +66,8 @@ async def test_mock_unstructred_response(tmp_path):
 
     # don't error on valid response
     adapter = MockAdapter(task, response={"setup": "asdf", "punchline": "asdf"})
-    answer = await adapter.invoke_returning_raw("You are a mock, send me the response!")
+    run = await adapter.invoke("You are a mock, send me the response!")
+    answer = json.loads(run.output.output)
     assert answer["setup"] == "asdf"
     assert answer["punchline"] == "asdf"
 
@@ -76,9 +77,12 @@ async def test_mock_unstructred_response(tmp_path):
         answer = await adapter.invoke("You are a mock, send me the response!")
 
     adapter = MockAdapter(task, response="string instead of dict")
-    with pytest.raises(RuntimeError):
+    with pytest.raises(
+        ValueError,
+        match="This task requires JSON output but the model didn't return valid JSON",
+    ):
         # Not a structed response so should error
-        answer = await adapter.invoke("You are a mock, send me the response!")
+        run = await adapter.invoke("You are a mock, send me the response!")
 
     # Should error, expecting a string, not a dict
     project = datamodel.Project(name="test", path=tmp_path / "test.kiln")
@@ -143,7 +147,8 @@ async def run_structured_output_test(tmp_path: Path, model_name: str, provider:
     task = build_structured_output_test_task(tmp_path)
     a = adapter_for_task(task, model_name=model_name, provider=provider)
     try:
-        parsed = await a.invoke_returning_raw("Cows")  # a joke about cows
+        run = await a.invoke("Cows")  # a joke about cows
+        parsed = json.loads(run.output.output)
     except ValueError as e:
         if str(e) == "Failed to connect to Ollama. Ensure Ollama is running.":
             pytest.skip(
@@ -162,6 +167,12 @@ async def run_structured_output_test(tmp_path: Path, model_name: str, provider:
         assert rating >= 0
         assert rating <= 10
 
+    # Check reasoning models
+    assert a._model_provider is not None
+    if a._model_provider.reasoning_capable:
+        assert "reasoning" in run.intermediate_outputs
+        assert isinstance(run.intermediate_outputs["reasoning"], str)
+
 
 def build_structured_input_test_task(tmp_path: Path):
     project = datamodel.Project(name="test", path=tmp_path / "test.kiln")
@@ -220,7 +231,8 @@ async def run_structured_input_task(
         await a.invoke({"a": 1, "b": 2, "d": 3})
 
     try:
-        response = await a.invoke_returning_raw({"a": 2, "b": 2, "c": 2})
+        run = await a.invoke({"a": 2, "b": 2, "c": 2})
+        response = run.output.output
     except ValueError as e:
         if str(e) == "Failed to connect to Ollama. Ensure Ollama is running.":
             pytest.skip(
@@ -241,6 +253,12 @@ async def run_structured_input_task(
     assert a.run_config.model_name == model_name
     assert a.run_config.model_provider_name == provider
 
+    # Check reasoning models
+    assert a._model_provider is not None
+    if a._model_provider.reasoning_capable:
+        assert "reasoning" in run.intermediate_outputs
+        assert isinstance(run.intermediate_outputs["reasoning"], str)
+
 
 @pytest.mark.paid
 async def test_structured_input_gpt_4o_mini(tmp_path):

kiln_ai/adapters/ollama_tools.py

@@ -38,6 +38,7 @@ async def ollama_online() -> bool:
 
 class OllamaConnection(BaseModel):
     message: str
+    version: str | None = None
     supported_models: List[str]
     untested_models: List[str] = Field(default_factory=list)
 
@@ -49,7 +50,7 @@ class OllamaConnection(BaseModel):
 def parse_ollama_tags(tags: Any) -> OllamaConnection | None:
     # Build a list of models we support for Ollama from the built-in model list
     supported_ollama_models = [
-        provider.provider_options["model"]
+        provider.model_id
         for model in built_in_models
         for provider in model.providers
         if provider.name == ModelProviderName.ollama
@@ -60,7 +61,7 @@ def parse_ollama_tags(tags: Any) -> OllamaConnection | None:
             alias
             for model in built_in_models
             for provider in model.providers
-            for alias in provider.provider_options.get("model_aliases", [])
+            for alias in provider.ollama_model_aliases or []
         ]
     )
 

kiln_ai/adapters/parsers/r1_parser.py

@@ -20,21 +20,33 @@ class R1ThinkingParser(BaseParser):
         Raises:
             ValueError: If response format is invalid (missing tags, multiple tags, or no content after closing tag)
         """
+
+        # The upstream providers (litellm, openrouter, fireworks) all keep changing their response formats, sometimes adding reasoning parsing where it didn't previously exist.
+        # If they do it already, great just return. If not we parse it ourselves. Not ideal, but better than upstream changes breaking the app.
+        if (
+            original_output.intermediate_outputs is not None
+            and "reasoning" in original_output.intermediate_outputs
+        ):
+            return original_output
+
         # This parser only works for strings
         if not isinstance(original_output.output, str):
             raise ValueError("Response must be a string for R1 parser")
 
         # Strip whitespace and validate basic structure
         cleaned_response = original_output.output.strip()
-        if not cleaned_response.startswith(self.START_TAG):
-            raise ValueError("Response must start with <think> tag")
 
         # Find the thinking tags
-        think_start = cleaned_response.find(self.START_TAG)
         think_end = cleaned_response.find(self.END_TAG)
+        if think_end == -1:
+            raise ValueError("Missing </think> tag")
 
-        if think_start == -1 or think_end == -1:
-            raise ValueError("Missing thinking tags")
+        think_tag_start = cleaned_response.find(self.START_TAG)
+        if think_tag_start == -1:
+            # We allow no start <think>, thinking starts on first char. QwQ does this.
+            think_start = 0
+        else:
+            think_start = think_tag_start + len(self.START_TAG)
 
         # Check for multiple tags
         if (
@@ -44,9 +56,7 @@ class R1ThinkingParser(BaseParser):
             raise ValueError("Multiple thinking tags found")
 
         # Extract thinking content
-        thinking_content = cleaned_response[
-            think_start + len(self.START_TAG) : think_end
-        ].strip()
+        thinking_content = cleaned_response[think_start:think_end].strip()
 
         # Extract result (everything after </think>)
         result = cleaned_response[think_end + len(self.END_TAG) :].strip()
@@ -54,16 +64,11 @@ class R1ThinkingParser(BaseParser):
         if not result or len(result) == 0:
             raise ValueError("No content found after </think> tag")
 
-        # Parse JSON if needed
-        output = result
-        if self.structured_output:
-            output = parse_json_string(result)
-
         # Add thinking content to intermediate outputs if it exists
         intermediate_outputs = original_output.intermediate_outputs or {}
         intermediate_outputs["reasoning"] = thinking_content
 
         return RunOutput(
-            output=output,
+            output=result,
             intermediate_outputs=intermediate_outputs,
         )

kiln_ai/adapters/parsers/test_r1_parser.py

@@ -19,6 +19,16 @@ def test_valid_response(parser):
     assert parsed.output == "This is the result"
 
 
+def test_already_parsed_response(parser):
+    response = RunOutput(
+        output="This is the result",
+        intermediate_outputs={"reasoning": "This is thinking content"},
+    )
+    parsed = parser.parse_output(response)
+    assert parsed.intermediate_outputs["reasoning"] == "This is thinking content"
+    assert parsed.output == "This is the result"
+
+
 def test_response_with_whitespace(parser):
     response = RunOutput(
         output="""
@@ -37,14 +47,16 @@ def test_response_with_whitespace(parser):
 
 
 def test_missing_start_tag(parser):
-    with pytest.raises(ValueError, match="Response must start with <think> tag"):
-        parser.parse_output(
-            RunOutput(output="Some content</think>result", intermediate_outputs=None)
-        )
+    parsed = parser.parse_output(
+        RunOutput(output="Some content</think>result", intermediate_outputs=None)
+    )
+
+    assert parsed.intermediate_outputs["reasoning"] == "Some content"
+    assert parsed.output == "result"
 
 
 def test_missing_end_tag(parser):
-    with pytest.raises(ValueError, match="Missing thinking tags"):
+    with pytest.raises(ValueError, match="Missing </think> tag"):
         parser.parse_output(
             RunOutput(output="<think>Some content", intermediate_outputs=None)
         )