kiln-ai 0.16.0__py3-none-any.whl → 0.18.0__py3-none-any.whl

kiln_ai/adapters/__init__.py

@@ -17,6 +17,7 @@ The eval submodule contains the code for evaluating the performance of a model.
 """
 
 from . import (
+    chat,
     data_gen,
     eval,
     fine_tune,
@@ -28,6 +29,7 @@ from . import (
 
 __all__ = [
     "model_adapters",
+    "chat",
     "data_gen",
     "fine_tune",
     "ml_model_list",

kiln_ai/adapters/adapter_registry.py

@@ -7,31 +7,33 @@ from kiln_ai.adapters.model_adapters.litellm_adapter import (
     LiteLlmAdapter,
     LiteLlmConfig,
 )
-from kiln_ai.adapters.provider_tools import core_provider, lite_llm_config
-from kiln_ai.datamodel import PromptId
+from kiln_ai.adapters.provider_tools import (
+    core_provider,
+    lite_llm_config_for_openai_compatible,
+)
+from kiln_ai.datamodel.task import RunConfigProperties
 from kiln_ai.utils.config import Config
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 
 def adapter_for_task(
     kiln_task: datamodel.Task,
-    model_name: str,
-    provider: ModelProviderName,
-    prompt_id: PromptId | None = None,
+    run_config_properties: RunConfigProperties,
     base_adapter_config: AdapterConfig | None = None,
 ) -> BaseAdapter:
     # Get the provider to run. For things like the fine-tune provider, we want to run the underlying provider
-    core_provider_name = core_provider(model_name, provider)
+    core_provider_name = core_provider(
+        run_config_properties.model_name, run_config_properties.model_provider_name
+    )
 
     match core_provider_name:
         case ModelProviderName.openrouter:
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
                 config=LiteLlmConfig(
-                    model_name=model_name,
+                    run_config_properties=run_config_properties,
                     base_url=getenv("OPENROUTER_BASE_URL")
                     or "https://openrouter.ai/api/v1",
-                    provider_name=provider,
                     default_headers={
                         "HTTP-Referer": "https://getkiln.ai/openrouter",
                         "X-Title": "KilnAI",
@@ -40,38 +42,32 @@ def adapter_for_task(
                         "api_key": Config.shared().open_router_api_key,
                     },
                 ),
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
             )
         case ModelProviderName.openai:
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
                 config=LiteLlmConfig(
-                    model_name=model_name,
-                    provider_name=provider,
+                    run_config_properties=run_config_properties,
                     additional_body_options={
                         "api_key": Config.shared().open_ai_api_key,
                     },
                 ),
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
             )
         case ModelProviderName.openai_compatible:
-            config = lite_llm_config(model_name)
+            config = lite_llm_config_for_openai_compatible(run_config_properties)
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
                 config=config,
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
             )
         case ModelProviderName.groq:
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
                 config=LiteLlmConfig(
-                    model_name=model_name,
-                    provider_name=provider,
+                    run_config_properties=run_config_properties,
                     additional_body_options={
                         "api_key": Config.shared().groq_api_key,
                     },
@@ -80,11 +76,9 @@ def adapter_for_task(
         case ModelProviderName.amazon_bedrock:
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
                 config=LiteLlmConfig(
-                    model_name=model_name,
-                    provider_name=provider,
+                    run_config_properties=run_config_properties,
                     additional_body_options={
                         "aws_access_key_id": Config.shared().bedrock_access_key,
                         "aws_secret_access_key": Config.shared().bedrock_secret_key,
@@ -99,11 +93,9 @@ def adapter_for_task(
             )
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
                 config=LiteLlmConfig(
-                    model_name=model_name,
-                    provider_name=provider,
+                    run_config_properties=run_config_properties,
                     # Set the Ollama base URL for 2 reasons:
                     # 1. To use the correct base URL
                     # 2. We use Ollama's OpenAI compatible API (/v1), and don't just let litellm use the Ollama API. We use more advanced features like json_schema.
@@ -117,11 +109,9 @@ def adapter_for_task(
         case ModelProviderName.fireworks_ai:
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
                 config=LiteLlmConfig(
-                    model_name=model_name,
-                    provider_name=provider,
+                    run_config_properties=run_config_properties,
                     additional_body_options={
                         "api_key": Config.shared().fireworks_api_key,
                     },
@@ -130,11 +120,9 @@ def adapter_for_task(
         case ModelProviderName.anthropic:
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
                 config=LiteLlmConfig(
-                    model_name=model_name,
-                    provider_name=provider,
+                    run_config_properties=run_config_properties,
                     additional_body_options={
                         "api_key": Config.shared().anthropic_api_key,
                     },
@@ -143,11 +131,9 @@ def adapter_for_task(
         case ModelProviderName.gemini_api:
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
                 config=LiteLlmConfig(
-                    model_name=model_name,
-                    provider_name=provider,
+                    run_config_properties=run_config_properties,
                     additional_body_options={
                         "api_key": Config.shared().gemini_api_key,
                     },
@@ -156,11 +142,9 @@ def adapter_for_task(
         case ModelProviderName.vertex:
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
                 config=LiteLlmConfig(
-                    model_name=model_name,
-                    provider_name=provider,
+                    run_config_properties=run_config_properties,
                     additional_body_options={
                         "vertex_project": Config.shared().vertex_project_id,
                         "vertex_location": Config.shared().vertex_location,
@@ -170,11 +154,9 @@ def adapter_for_task(
         case ModelProviderName.together_ai:
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
                 config=LiteLlmConfig(
-                    model_name=model_name,
-                    provider_name=provider,
+                    run_config_properties=run_config_properties,
                     additional_body_options={
                         "api_key": Config.shared().together_api_key,
                     },
@@ -183,12 +165,10 @@ def adapter_for_task(
         case ModelProviderName.azure_openai:
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
                 config=LiteLlmConfig(
                     base_url=Config.shared().azure_openai_endpoint,
-                    model_name=model_name,
-                    provider_name=provider,
+                    run_config_properties=run_config_properties,
                     additional_body_options={
                         "api_key": Config.shared().azure_openai_api_key,
                         "api_version": "2025-02-01-preview",
@@ -198,11 +178,9 @@ def adapter_for_task(
         case ModelProviderName.huggingface:
             return LiteLlmAdapter(
                 kiln_task=kiln_task,
-                prompt_id=prompt_id,
                 base_adapter_config=base_adapter_config,
                 config=LiteLlmConfig(
-                    model_name=model_name,
-                    provider_name=provider,
+                    run_config_properties=run_config_properties,
                     additional_body_options={
                         "api_key": Config.shared().huggingface_api_key,
                     },

kiln_ai/adapters/chat/__init__.py

@@ -0,0 +1,8 @@
+from .chat_formatter import (
+    ChatFormatter,
+    ChatMessage,
+    ChatStrategy,
+    get_chat_formatter,
+)
+
+__all__ = ["ChatFormatter", "ChatMessage", "ChatStrategy", "get_chat_formatter"]

kiln_ai/adapters/chat/chat_formatter.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+import json
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Dict, List, Literal, Optional
+
+from kiln_ai.datamodel.datamodel_enums import ChatStrategy
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+COT_FINAL_ANSWER_PROMPT = "Considering the above, return a final result."
+
+
+@dataclass
+class ChatMessage:
+    role: Literal["system", "assistant", "user"]
+    content: Optional[str]
+
+
+@dataclass
+class ChatTurn:
+    """
+    All data needed to send a chat turn to the model.
+    """
+
+    messages: List[ChatMessage]
+    final_call: bool
+
+
+class ChatFormatter(ABC):
+    def __init__(
+        self,
+        system_message: str,
+        user_input: str | Dict,
+        thinking_instructions: str | None = None,
+    ) -> None:
+        self.system_message = system_message
+        self.user_input = user_input
+        self.thinking_instructions = thinking_instructions
+        self._messages: List[ChatMessage] = []
+        self._state = "start"
+        self._intermediate_outputs: Dict[str, str] = {}
+
+    @property
+    def messages(self) -> List[ChatMessage]:
+        return list(self._messages)
+
+    def message_dicts(self) -> List[dict[str, str | None]]:
+        return [{"role": m.role, "content": m.content} for m in self._messages]
+
+    def intermediate_outputs(self) -> Dict[str, str]:
+        """Get the intermediate outputs from the chat formatter."""
+        return self._intermediate_outputs
+
+    @abstractmethod
+    def next_turn(self, previous_output: str | None = None) -> Optional[ChatTurn]:
+        """Advance the conversation and return the next messages if any."""
+        raise NotImplementedError
+
+
+class SingleTurnFormatter(ChatFormatter):
+    def next_turn(self, previous_output: str | None = None) -> Optional[ChatTurn]:
+        if self._state == "start":
+            msgs = [
+                ChatMessage("system", self.system_message),
+                ChatMessage("user", format_user_message(self.user_input)),
+            ]
+            self._state = "awaiting_final"
+            self._messages.extend(msgs)
+            return ChatTurn(messages=msgs, final_call=True)
+
+        if self._state == "awaiting_final":
+            if previous_output is None:
+                raise ValueError("previous_output required for final step")
+            self._messages.append(ChatMessage("assistant", previous_output))
+            self._state = "done"
+            return None
+
+        return None
+
+
+class TwoMessageCotLegacyFormatter(ChatFormatter):
+    def __init__(
+        self,
+        system_message: str,
+        user_input: str | Dict,
+        thinking_instructions: str | None,
+    ) -> None:
+        super().__init__(system_message, user_input, thinking_instructions)
+        if self.thinking_instructions is None:
+            raise ValueError(
+                "thinking_instructions are required when strategy is final_and_intermediate"
+            )
+
+    def next_turn(self, previous_output: str | None = None) -> Optional[ChatTurn]:
+        if self._state == "start":
+            msgs = [
+                ChatMessage("system", self.system_message),
+                ChatMessage("user", format_user_message(self.user_input)),
+                ChatMessage("system", self.thinking_instructions),
+            ]
+            self._state = "awaiting_thinking"
+            self._messages.extend(msgs)
+            return ChatTurn(messages=msgs, final_call=False)
+
+        if self._state == "awaiting_thinking":
+            if previous_output is None:
+                raise ValueError("previous_output required for thinking step")
+            msgs = [
+                ChatMessage("assistant", previous_output),
+                ChatMessage("user", COT_FINAL_ANSWER_PROMPT),
+            ]
+            self._intermediate_outputs["chain_of_thought"] = previous_output
+            self._state = "awaiting_final"
+            self._messages.extend(msgs)
+            return ChatTurn(messages=msgs, final_call=True)
+
+        if self._state == "awaiting_final":
+            if previous_output is None:
+                raise ValueError("previous_output required for final step")
+            self._messages.append(ChatMessage("assistant", previous_output))
+            self._state = "done"
+            return None
+
+        return None
+
+
+class TwoMessageCotFormatter(ChatFormatter):
+    def __init__(
+        self,
+        system_message: str,
+        user_input: str | Dict,
+        thinking_instructions: str | None,
+    ) -> None:
+        super().__init__(system_message, user_input, thinking_instructions)
+        if self.thinking_instructions is None:
+            raise ValueError(
+                "thinking_instructions are required when strategy is final_and_intermediate"
+            )
+
+    def next_turn(self, previous_output: str | None = None) -> Optional[ChatTurn]:
+        if self._state == "start":
+            # User message combines the input and the thinking instructions
+            formatted_user_message = format_user_message(self.user_input)
+            user_message = f"The input is:\n<user_input>\n{formatted_user_message}\n</user_input>\n\n{self.thinking_instructions}"
+
+            msgs = [
+                ChatMessage("system", self.system_message),
+                ChatMessage("user", user_message),
+            ]
+            self._state = "awaiting_thinking"
+            self._messages.extend(msgs)
+            return ChatTurn(messages=msgs, final_call=False)
+
+        if self._state == "awaiting_thinking":
+            if previous_output is None:
+                raise ValueError("previous_output required for thinking step")
+            msgs = [
+                ChatMessage("assistant", previous_output),
+                ChatMessage("user", COT_FINAL_ANSWER_PROMPT),
+            ]
+            self._intermediate_outputs["chain_of_thought"] = previous_output
+            self._state = "awaiting_final"
+            self._messages.extend(msgs)
+            return ChatTurn(messages=msgs, final_call=True)
+
+        if self._state == "awaiting_final":
+            if previous_output is None:
+                raise ValueError("previous_output required for final step")
+            self._messages.append(ChatMessage("assistant", previous_output))
+            self._state = "done"
+            return None
+
+        return None
+
+
+class SingleTurnR1ThinkingFormatter(ChatFormatter):
+    def next_turn(self, previous_output: str | None = None) -> Optional[ChatTurn]:
+        if self._state == "start":
+            msgs = [
+                ChatMessage("system", self.system_message),
+                ChatMessage("user", format_user_message(self.user_input)),
+            ]
+            self._state = "awaiting_final"
+            self._messages.extend(msgs)
+            return ChatTurn(messages=msgs, final_call=True)
+
+        if self._state == "awaiting_final":
+            if previous_output is None:
+                raise ValueError("previous_output required for final step")
+            self._messages.append(ChatMessage("assistant", previous_output))
+            self._state = "done"
+            return None
+
+        return None
+
+
+def get_chat_formatter(
+    strategy: ChatStrategy,
+    system_message: str,
+    user_input: str | Dict,
+    thinking_instructions: str | None = None,
+) -> ChatFormatter:
+    match strategy:
+        case ChatStrategy.single_turn:
+            return SingleTurnFormatter(system_message, user_input)
+        case ChatStrategy.two_message_cot_legacy:
+            return TwoMessageCotLegacyFormatter(
+                system_message, user_input, thinking_instructions
+            )
+        case ChatStrategy.two_message_cot:
+            return TwoMessageCotFormatter(
+                system_message, user_input, thinking_instructions
+            )
+        case ChatStrategy.single_turn_r1_thinking:
+            return SingleTurnR1ThinkingFormatter(system_message, user_input)
+        case _:
+            raise_exhaustive_enum_error(strategy)
+
+
+def format_user_message(input: Dict | str) -> str:
+    """Build a user message from the input.
+
+    Args:
+        input (Union[Dict, str]): The input to format into a message.
+
+    Returns:
+        str: The formatted user message.
+    """
+    if isinstance(input, dict):
+        return json.dumps(input, ensure_ascii=False)
+
+    return input

kiln_ai/adapters/chat/test_chat_formatter.py

@@ -0,0 +1,131 @@
+from kiln_ai.adapters.chat import ChatStrategy, get_chat_formatter
+from kiln_ai.adapters.chat.chat_formatter import (
+    COT_FINAL_ANSWER_PROMPT,
+    format_user_message,
+)
+
+
+def test_chat_formatter_final_only():
+    expected = [
+        {"role": "system", "content": "system message"},
+        {"role": "user", "content": "test input"},
+        {"role": "assistant", "content": "test output"},
+    ]
+
+    formatter = get_chat_formatter(
+        strategy=ChatStrategy.single_turn,
+        system_message="system message",
+        user_input="test input",
+    )
+
+    first = formatter.next_turn()
+    assert [m.__dict__ for m in first.messages] == expected[:2]
+    assert first.final_call
+    assert formatter.intermediate_outputs() == {}
+
+    assert formatter.next_turn("test output") is None
+    assert formatter.message_dicts() == expected
+    assert formatter.intermediate_outputs() == {}
+
+
+def test_chat_formatter_final_and_intermediate():
+    expected = [
+        {"role": "system", "content": "system message"},
+        {"role": "user", "content": "test input"},
+        {"role": "system", "content": "thinking instructions"},
+        {"role": "assistant", "content": "thinking output"},
+        {"role": "user", "content": COT_FINAL_ANSWER_PROMPT},
+        {"role": "assistant", "content": "test output"},
+    ]
+
+    formatter = get_chat_formatter(
+        strategy=ChatStrategy.two_message_cot_legacy,
+        system_message="system message",
+        user_input="test input",
+        thinking_instructions="thinking instructions",
+    )
+
+    first = formatter.next_turn()
+    assert [m.__dict__ for m in first.messages] == expected[:3]
+    assert not first.final_call
+    assert formatter.intermediate_outputs() == {}
+
+    second = formatter.next_turn("thinking output")
+    assert [m.__dict__ for m in second.messages] == expected[3:5]
+    assert second.final_call
+    assert formatter.intermediate_outputs() == {"chain_of_thought": "thinking output"}
+
+    assert formatter.next_turn("test output") is None
+    assert formatter.message_dicts() == expected
+    assert formatter.intermediate_outputs() == {"chain_of_thought": "thinking output"}
+
+
+def test_chat_formatter_two_message_cot():
+    user_message = "The input is:\n<user_input>\ntest input\n</user_input>\n\nthinking instructions"
+    expected = [
+        {"role": "system", "content": "system message"},
+        {"role": "user", "content": user_message},
+        {"role": "assistant", "content": "thinking output"},
+        {"role": "user", "content": COT_FINAL_ANSWER_PROMPT},
+        {"role": "assistant", "content": "test output"},
+    ]
+
+    formatter = get_chat_formatter(
+        strategy=ChatStrategy.two_message_cot,
+        system_message="system message",
+        user_input="test input",
+        thinking_instructions="thinking instructions",
+    )
+
+    first = formatter.next_turn()
+    assert [m.__dict__ for m in first.messages] == expected[:2]
+    assert not first.final_call
+    assert formatter.intermediate_outputs() == {}
+
+    second = formatter.next_turn("thinking output")
+    assert [m.__dict__ for m in second.messages] == expected[2:4]
+    assert second.final_call
+    assert formatter.intermediate_outputs() == {"chain_of_thought": "thinking output"}
+
+    assert formatter.next_turn("test output") is None
+    assert formatter.message_dicts() == expected
+    assert formatter.intermediate_outputs() == {"chain_of_thought": "thinking output"}
+
+
+def test_chat_formatter_r1_style():
+    thinking_output = "<think>thinking</think> answer"
+    expected = [
+        {"role": "system", "content": "system message"},
+        {"role": "user", "content": "test input"},
+        {"role": "assistant", "content": thinking_output},
+    ]
+
+    formatter = get_chat_formatter(
+        strategy=ChatStrategy.single_turn_r1_thinking,
+        system_message="system message",
+        user_input="test input",
+    )
+
+    first = formatter.next_turn()
+    assert [m.__dict__ for m in first.messages] == expected[:2]
+    assert first.final_call
+
+    assert formatter.next_turn(thinking_output) is None
+    assert formatter.message_dicts() == expected
+    assert formatter.intermediate_outputs() == {}
+
+
+def test_format_user_message():
+    # String
+    assert format_user_message("test input") == "test input"
+    # JSON, preserving order
+    assert (
+        format_user_message({"test": "input", "a": "b"})
+        == '{"test": "input", "a": "b"}'
+    )
+
+
+def test_simple_prompt_builder_structured_input_non_ascii():
+    input = {"key": "你好👋"}
+    user_msg = format_user_message(input)
+    assert "你好👋" in user_msg